tenstorrent / tt-metal

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

ttnn implementation for Swin_S model #8590

Open saichandax opened 4 months ago

saichandax commented 4 months ago

reference model is from: https://github.com/pytorch/vision/blob/main/torchvision/models/swin_transformer.py

HariniMohan0102 commented 4 months ago

To implement Swin_s transformer model, model is segregated into submodules where each module is dependent on the other. The sequence of submodules are

  1. ShiftedWindowAttention
  2. SwinTransformerBlock
  3. PatchMerging
  4. SwinTransformer

Currently working on ShiftedWindowAttention submodule. Task Updates regarding ShiftedWindowAttention:

For your reference, added torch and referenc graphs of ShiftedWindowAttention submodule. Torch graph: torch_graph

Reference graph: ref_graph