File size: 5,190 Bytes
6fd6195 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 |
# Copyright 2020 - 2022 MONAI Consortium
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# http://www.apache.org/licenses/LICENSE-2.0
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import torch
import torch.nn as nn
from monai.networks.nets.swin_unetr import SwinTransformer as SwinViT
from monai.utils import ensure_tuple_rep
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) # YY [2,2,2]
window_size = ensure_tuple_rep(7, args.spatial_dims) # YY [7,7,7]
dim = args.bottleneck_depth # YY 768 or 1536
self.swinViT = SwinViT(
in_chans=args.in_channels, # 2 for T1T2
embed_dim=args.feature_size, # YY 48, 96
window_size=window_size,
patch_size=patch_size,
depths = args.num_swin_blocks_per_stage, # YY [2, 2, 2, 2],
num_heads = args.num_heads_per_stage, # YY [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,
)
self.rotation_pre = nn.Identity()
self.rotation_head = nn.Linear(dim, 4) # YY rotation 4 value
self.contrastive_pre = nn.Identity()
self.contrastive_head = nn.Linear(dim, 512) # YY why 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),
)
def forward(self, x):
x_out = self.swinViT(x.contiguous())[4] # YY why [4]. [x0_out, x1_out, x2_out, x3_out, x4_out] stage4 meaning that get the output of final stage in Encoder
_, c, h, w, d = x_out.shape # # [4, 768, 3, 3, 3]
x4_reshape = x_out.flatten(start_dim=2, end_dim=4) # # [4, 768, 27]
x4_reshape = x4_reshape.transpose(1, 2) # [4, 27, 768]
x_rot = self.rotation_pre(x4_reshape[:, 0]) # [4, 768]
x_rot = self.rotation_head(x_rot) # [4, 4]
x_contrastive = self.contrastive_pre(x4_reshape[:, 1]) # [4, 768]
x_contrastive = self.contrastive_head(x_contrastive) # [4, 512]
x_rec = x_out.flatten(start_dim=2, end_dim=4) # [4, 768, 27]
x_rec = x_rec.view(-1, c, h, w, d) # [4, 768, 3, 3, 3]
x_rec = self.conv(x_rec) # # [4, in_channel, 96, 96, 96]
return x_rot, x_contrastive, x_rec
|