Transformer-Unet

Unet implementation with transformer encoder

Usage

from unet import TransUnet

a = torch.rand(2, 3, 128, 128)

model = TransUnet(in_channels=3, img_dim=128, vit_blocks=1, vit_dim_linear_mhsa_block=512, classes=5)
y = model(a)
print('final out shape:', y.shape)