I am trying to use the swin unetr pretraining pipeline to train a model on public datasets.
after the training finished (I did not changed anything in the config file or data loader and just used the code as it is), I want to visualize some the attention map on the CT images.
class SSLHead(nn.Module):
def __init__(self, args, upsample="vae", dim=768):
super(SSLHead, self).__init__()
patch_size = ensure_tuple_rep(2, args.spatial_dims)
window_size = ensure_tuple_rep(7, args.spatial_dims)
self.swinViT = SwinViT(
in_chans=args.in_channels,
embed_dim=args.feature_size,
window_size=window_size,
patch_size=patch_size,
depths=[2, 2, 2, 2],
num_heads=[3, 6, 12, 24],
mlp_ratio=4.0,
qkv_bias=True,
drop_rate=0.0,
attn_drop_rate=0.0,
drop_path_rate=args.dropout_path_rate,
norm_layer=torch.nn.LayerNorm,
use_checkpoint=args.use_checkpoint,
spatial_dims=args.spatial_dims,
use_v2 =args.use_v2,
)
self.rotation_pre = nn.Identity()
self.rotation_head = nn.Linear(dim, 4)
self.contrastive_pre = nn.Identity()
self.contrastive_head = nn.Linear(dim, 512)
if upsample == "large_kernel_deconv":
self.conv = nn.ConvTranspose3d(dim, args.in_channels, kernel_size=(32, 32, 32), stride=(32, 32, 32))
elif upsample == "deconv":
self.conv = nn.Sequential(
nn.ConvTranspose3d(dim, dim // 2, kernel_size=(2, 2, 2), stride=(2, 2, 2)),
nn.ConvTranspose3d(dim // 2, dim // 4, kernel_size=(2, 2, 2), stride=(2, 2, 2)),
nn.ConvTranspose3d(dim // 4, dim // 8, kernel_size=(2, 2, 2), stride=(2, 2, 2)),
nn.ConvTranspose3d(dim // 8, dim // 16, kernel_size=(2, 2, 2), stride=(2, 2, 2)),
nn.ConvTranspose3d(dim // 16, args.in_channels, kernel_size=(2, 2, 2), stride=(2, 2, 2)),
)
elif upsample == "vae":
self.conv = nn.Sequential(
nn.Conv3d(dim, dim // 2, kernel_size=3, stride=1, padding=1),
nn.InstanceNorm3d(dim // 2),
nn.LeakyReLU(),
nn.Upsample(scale_factor=2, mode="trilinear", align_corners=False),
nn.Conv3d(dim // 2, dim // 4, kernel_size=3, stride=1, padding=1),
nn.InstanceNorm3d(dim // 4),
nn.LeakyReLU(),
nn.Upsample(scale_factor=2, mode="trilinear", align_corners=False),
nn.Conv3d(dim // 4, dim // 8, kernel_size=3, stride=1, padding=1),
nn.InstanceNorm3d(dim // 8),
nn.LeakyReLU(),
nn.Upsample(scale_factor=2, mode="trilinear", align_corners=False),
nn.Conv3d(dim // 8, dim // 16, kernel_size=3, stride=1, padding=1),
nn.InstanceNorm3d(dim // 16),
nn.LeakyReLU(),
nn.Upsample(scale_factor=2, mode="trilinear", align_corners=False),
nn.Conv3d(dim // 16, dim // 16, kernel_size=3, stride=1, padding=1),
nn.InstanceNorm3d(dim // 16),
nn.LeakyReLU(),
nn.Upsample(scale_factor=2, mode="trilinear", align_corners=False),
nn.Conv3d(dim // 16, args.in_channels, kernel_size=1, stride=1),
)
@torch.no_grad()
def extract_grid(self, x):
with torch.cuda.amp.autocast():
return self.swinViT(x.contiguous())[4]
@torch.no_grad()
def extract_global(self, x, pool="mean"):
g = self.extract_grid(x) # (B, 768, Z, Y, X)
if pool == "mean":
return g.mean(dim=(2, 3, 4))
elif pool == "max":
return g.max(dim=2)[0].max(dim=2)[0].max(dim=2)[0]
else:
raise ValueError(pool)
def forward(self, x):
x_out = self.swinViT(x.contiguous())[4]
_, c, h, w, d = x_out.shape
x4_reshape = x_out.flatten(start_dim=2, end_dim=4)
x4_reshape = x4_reshape.transpose(1, 2)
x_rot = self.rotation_pre(x4_reshape[:, 0])
x_rot = self.rotation_head(x_rot)
x_contrastive = self.contrastive_pre(x4_reshape[:, 1])
x_contrastive = self.contrastive_head(x_contrastive)
x_rec = x_out.flatten(start_dim=2, end_dim=4)
x_rec = x_rec.view(-1, c, h, w, d)
x_rec = self.conv(x_rec)
return x_rot, x_contrastive, x_rec
class SwinEncoder(nn.Module):
def __init__(self,
ssl_head: SSLHead,
pooled: bool = False,
pool_type: str = "mean"):
super().__init__()
self.backbone = ssl_head.swinViT
self.pooled = pooled
self.pool_type = pool_type.lower()
@torch.no_grad()
def forward(self, x, pooled: bool | None = None):
if pooled is None:
pooled = self.pooled
g = self.backbone(x.contiguous())[4] # (B, 768, Z, Y, X)
if not pooled:
return g
if self.pool_type == "mean":
return g.mean(dim=(2, 3, 4)) # (B, 768)
elif self.pool_type == "max":
return g.amax(dim=(2, 3, 4)) # (B, 768)
else:
raise ValueError(f"Unknown pool_type {self.pool_type!r}")
def get_loader(args):
jsonlist = args.json_path
num_workers = args.num_workers
val_list = load_decathlon_datalist(jsonlist, False, "validation")
print("number of validation data: {}".format(len(val_list)))
tst_transforms = Compose(
[
LoadImaged(keys=["image"], ensure_channel_first=True, image_only=False),
Orientationd(keys=["image"], axcodes="RAS"),
ScaleIntensityRanged(
keys=["image"], a_min=args.a_min, a_max=args.a_max, b_min=args.b_min, b_max=args.b_max, clip=True
),
CropForegroundd(keys=["image"], source_key="image", k_divisible=[args.roi_x, args.roi_y, args.roi_z], allow_smaller = True),
SpatialPadd(keys="image", spatial_size=[args.roi_x, args.roi_y, args.roi_z]),
ToTensord(keys=["image"]),
]
)
print("Using generic dataset")
tst_ds = Dataset(data=val_list, transform=tst_transforms)
tst_loader = DataLoader(
tst_ds, batch_size=args.batch_size, num_workers=num_workers, shuffle=False, drop_last=True
)
return tst_loader
root_save = 'path/to/the/code'
json_path = os.path.join(root_code, 'Pys', 'jsons', 'SwinUNETRPretaining.json')
config_path = os.path.join(root_save, 'config.yaml')
model_pth = os.path.join(root_save, 'logs', 'model_bestValRMSE.pt')
with open(config_path, "r") as file:
cfg_dict = yaml.safe_load(file)
cfg = Namespace(**cfg_dict)
cfg.json_path = json_path
test_loader = get_loader(cfg)
ssl = SSLHead(cfg).cuda().eval()
ckpt = torch.load(model_pth, map_location="cuda")
state = ckpt["state_dict"]
state = OrderedDict((k.replace("module.", "", 1), v) for k, v in state.items())
ssl.load_state_dict(state, strict=True)
encoder = SwinEncoder(ssl, pooled=False).cuda().eval()
for step, batch in enumerate(test_loader):
x = batch['image'].cuda()
grid_embed = sliding_window_inference(inputs=x,
roi_size=(cfg.roi_x, cfg.roi_y, cfg.roi_z),
sw_batch_size=4,
predictor=encoder,
overlap=0.25, mode="constant", padding_mode="constant")
global_embed = grid_embed.mean(dim=(2,3,4))
print(f'input shape {x.shape}') #input shape torch.Size([1, 1, 512, 512, 128])
print(f'grid_embed shape {grid_embed.shape}') #grid_embed shape torch.Size([1, 768, 16, 16, 4])
print(f'global_embed shape {global_embed.shape}') #global_embed shape torch.Size([1, 768])
break
act_map = torch.linalg.norm(grid_embed[0], dim=0)
heat = act_map.unsqueeze(0).unsqueeze(0) # (1,1,16,16,4)
heat_up = F.interpolate(
heat, size=tuple(x.shape[2:]), # (D, H, W) = (128,512,512)
mode="trilinear",
align_corners=True
).squeeze()
heat_np = heat_up.cpu().numpy()
heat_np = (heat_np - heat_np.min()) / (heat_np.max() - heat_np.min() + 1e-8) #(512, 512, 128)
z = 64
plt.figure(figsize=(10,4))
plt.imshow(x[0,0,:, :, z].cpu(), cmap="gray")
plt.imshow(heat_np[:, :, z], cmap="hot", alpha=0.5)
plt.show()Any idea how can I get the actual attention map? I am looking for something like this:
Hello all,
I am trying to use the swin unetr pretraining pipeline to train a model on public datasets.
after the training finished (I did not changed anything in the config file or data loader and just used the code as it is), I want to visualize some the attention map on the CT images.
Any idea how can I get the actual attention map? I am looking for something like this:
Thanks a lot 🙂