tenstorrent / tt-metal

:metal: TT-NN operator library, and TT-Metalium low level kernel programming model.
Apache License 2.0
452 stars 65 forks source link

Vanilla UNet - Model card #13272

Closed punithsekar closed 3 weeks ago

punithsekar commented 3 weeks ago

Reference link:- https://github.com/mateuszbuda/brain-segmentation-pytorch/tree/d45f8908ab2f0246ba204c702a6161c9eb25f902

Vanilla Unet torch implementation of our reference implementation - https://github.com/tenstorrent/tt-metal/blob/05a8838c0d935d41112e07bc867b4cfc40115d09/models/experimental/functional_vanilla_unet/reference/unet.py

Our pytorch model graph - model_vanilla_unet.gv.pdf

Modules:

  1. UNet

Ops:

  1. Conv2d
  2. BatchNorm2d
  3. ReLU
  4. MaxPool2d
  5. cat
  6. ConvTranspose2d
  7. Sigmoid
punithsekar commented 3 weeks ago

fyi @saichandax

punithsekar commented 3 weeks ago

Closing this for now, We can open it if needed!