Spaces:
Runtime error
Runtime error
| from Models.BackBone import * | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| from einops import rearrange, repeat | |
| class DoubleConv(nn.Module): | |
| """(convolution => [BN] => ReLU) * 2""" | |
| def __init__(self, in_channels, out_channels, mid_channels=None): | |
| super(DoubleConv, self).__init__() | |
| if not mid_channels: | |
| mid_channels = out_channels | |
| self.double_conv = nn.Sequential( | |
| nn.Conv2d(in_channels, mid_channels, kernel_size=3, padding=1), | |
| nn.BatchNorm2d(mid_channels), | |
| nn.LeakyReLU(inplace=True), | |
| nn.Conv2d(mid_channels, out_channels, kernel_size=3, padding=1), | |
| nn.BatchNorm2d(out_channels), | |
| nn.LeakyReLU(inplace=True) | |
| ) | |
| def forward(self, x): | |
| return self.double_conv(x) | |
| class BoTMultiHeadAttention(nn.Module): | |
| def __init__(self, in_feature_dim, num_heads=8, dim_head=None, dropout_rate=0.): | |
| super().__init__() | |
| self.num_heads = num_heads | |
| self.dim_head = dim_head or in_feature_dim // num_heads | |
| self.scale = self.dim_head ** -0.5 | |
| inner_dim = self.dim_head * self.num_heads | |
| self.weights_qkv = nn.ModuleList([ | |
| nn.Linear(in_feature_dim, inner_dim, bias=False), | |
| nn.Linear(in_feature_dim, inner_dim, bias=False), | |
| nn.Linear(in_feature_dim, inner_dim, bias=False) | |
| ]) | |
| self.out_layer = nn.Sequential( | |
| nn.Linear(inner_dim, in_feature_dim), | |
| nn.Dropout(dropout_rate) | |
| ) | |
| self.layer_norm = nn.LayerNorm(in_feature_dim) | |
| def forward(self, q_s, k_s=None, v_s=None, pos_emb=None): | |
| if k_s is None and v_s is None: | |
| k_s = v_s = q_s | |
| elif v_s is None: | |
| v_s = k_s | |
| q, k, v = [self.weights_qkv[idx](x) for idx, x in enumerate([q_s, k_s, v_s])] | |
| q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h=self.num_heads), [q, k, v]) | |
| content_content_att = torch.einsum('b h i d, b h j d -> b h i j', q, k) * self.scale | |
| if pos_emb is not None: | |
| pos_emb = rearrange(pos_emb, 'b n (h d) -> b h n d', h=self.num_heads) | |
| content_position_att = torch.einsum('b h i d, b h j d -> b h i j', q, pos_emb) * self.scale | |
| att_mat = content_content_att + content_position_att | |
| else: | |
| att_mat = content_content_att | |
| att_mat = att_mat.softmax(dim=-1) | |
| atted_x = torch.einsum('b h i j , b h j d -> b h i d', att_mat, v) | |
| atted_x = rearrange(atted_x, 'b h n d -> b n (h d)') | |
| atted_x = self.out_layer(atted_x) | |
| out = self.layer_norm(atted_x + q_s) | |
| return out | |
| class STTNet(nn.Module): | |
| def __init__(self, in_channel, n_classes, *args, **kwargs): | |
| super(STTNet, self).__init__() | |
| self.in_channel = in_channel | |
| self.n_classes = n_classes | |
| # kwargs['backbone'] = res18, res50 or vgg16 | |
| self.res_backbone = get_backbone( | |
| model_name=kwargs['backbone'], num_classes=None, **kwargs | |
| ) | |
| # kwargs['out_keys'] = ['block_4'] or ['block_5'] | |
| self.last_block = kwargs['out_keys'][-1] | |
| if '18' in kwargs['backbone']: | |
| # 512 256 128 64 32 16 | |
| layer_channels = [64, 64, 128, 256, 512] | |
| self.reduce_dim_in = 256 | |
| self.reduce_dim_out = 256 // 4 | |
| elif '50' in kwargs['backbone']: | |
| layer_channels = [64, 256, 512, 1024, 2048] | |
| self.reduce_dim_in = 1024 | |
| self.reduce_dim_out = 1024 // 16 | |
| elif '16' in kwargs['backbone']: | |
| layer_channels = [64, 128, 256, 512, 512] | |
| self.reduce_dim_in = 512 | |
| self.reduce_dim_out = 512 // 8 | |
| self.f_map_size = 32 | |
| # kwargs['top_k_s'] = 64 | |
| self.top_k_s = kwargs['top_k_s'] | |
| # kwargs['top_k_c'] = 16 | |
| self.top_k_c = kwargs['top_k_c'] | |
| # kwargs['encoder_pos'] = True or False | |
| self.encoder_pos = kwargs['encoder_pos'] | |
| # kwargs['decoder_pos'] = True or False | |
| self.decoder_pos = kwargs['decoder_pos'] | |
| # kwargs['model_pattern'] = ['X', 'A', 'S', 'C'] means different features concatenation | |
| self.model_pattern = kwargs['model_pattern'] | |
| self.cat_num = len(self.model_pattern) | |
| if 'A' in self.model_pattern: | |
| self.cat_num += 1 | |
| self.num_head_s = max(2, min(self.top_k_s // 8, 64)) | |
| self.num_head_c = min(2, min(self.top_k_c // 4, 64)) | |
| self.reduce_channel_b5 = nn.Sequential( | |
| nn.Conv2d(in_channels=self.reduce_dim_in, out_channels=self.reduce_dim_out, kernel_size=1), | |
| nn.BatchNorm2d(self.reduce_dim_out), | |
| nn.LeakyReLU() | |
| ) | |
| # position embedding | |
| # if self.encoder_pos or self.decoder_pos: | |
| self.spatial_embedding_h = nn.Parameter( | |
| torch.randn(1, self.reduce_dim_out, self.f_map_size, 1), requires_grad=True) | |
| self.spatial_embedding_w = nn.Parameter( | |
| torch.randn(1, self.reduce_dim_out, 1, self.f_map_size), requires_grad=True) | |
| self.channel_embedding = nn.Parameter( | |
| torch.randn(1, self.reduce_dim_out, self.f_map_size ** 2), requires_grad=True) | |
| # spatial attention ops | |
| self.get_s_probability = nn.Sequential( | |
| nn.Conv2d(self.reduce_dim_out, self.reduce_dim_out // 4, kernel_size=3, padding=1), | |
| nn.BatchNorm2d(self.reduce_dim_out // 4), | |
| nn.LeakyReLU(inplace=True), | |
| nn.Conv2d(self.reduce_dim_out // 4, 1, kernel_size=3, padding=1), | |
| nn.Sigmoid() | |
| ) | |
| # b5 spatial encoder and decoder | |
| self.tf_encoder_spatial_b5 = BoTMultiHeadAttention( | |
| in_feature_dim=self.reduce_dim_out, | |
| num_heads=self.num_head_s | |
| ) | |
| self.tf_decoder_spatial_b5 = BoTMultiHeadAttention( | |
| in_feature_dim=self.reduce_dim_out, | |
| num_heads=self.num_head_s | |
| ) | |
| # channel attention ops | |
| self.get_c_probability = nn.Sequential( | |
| nn.Conv2d(self.reduce_dim_out, self.reduce_dim_out // 8, kernel_size=self.f_map_size), | |
| nn.BatchNorm2d(self.reduce_dim_out // 8), | |
| nn.LeakyReLU(inplace=True), | |
| nn.Conv2d(self.reduce_dim_out // 8, self.reduce_dim_out, kernel_size=1), | |
| nn.Sigmoid() | |
| ) | |
| # b5 channel encoder and decoder | |
| self.tf_encoder_channel_b5 = BoTMultiHeadAttention( | |
| in_feature_dim=self.f_map_size ** 2, | |
| num_heads=self.num_head_c | |
| ) | |
| self.tf_decoder_channel_b5 = BoTMultiHeadAttention( | |
| in_feature_dim=self.f_map_size ** 2, | |
| num_heads=self.num_head_c | |
| ) | |
| self.before_predict_head_conv = nn.Sequential( | |
| nn.Conv2d(in_channels=self.reduce_dim_out * self.cat_num, out_channels=self.reduce_dim_in, kernel_size=1), | |
| nn.BatchNorm2d(self.reduce_dim_in), | |
| nn.LeakyReLU() | |
| ) | |
| if self.last_block == 'block5': | |
| self.pre_pixel_shuffle = nn.PixelShuffle(2) | |
| # 128, 256, 256 | |
| self.pre_double_conv = DoubleConv( | |
| in_channels=layer_channels[4] // 4, | |
| out_channels=layer_channels[3], | |
| mid_channels=layer_channels[3] | |
| ) | |
| self.pixel_shuffle1 = nn.PixelShuffle(4) | |
| # 16, 64, 64 | |
| self.double_conv1 = DoubleConv( | |
| in_channels=layer_channels[3] // 16, | |
| out_channels=layer_channels[1], | |
| mid_channels=layer_channels[3] // 4 | |
| ) | |
| # 4, 16, 16 | |
| self.pixel_shuffle2 = nn.PixelShuffle(4) | |
| self.double_conv2 = DoubleConv( | |
| in_channels=layer_channels[1] // 16, | |
| out_channels=layer_channels[1] // 4, | |
| mid_channels=layer_channels[1] // 4 | |
| ) | |
| last_channels = layer_channels[1] // 4 | |
| # 16, 32 | |
| # 32, 2 | |
| if '18' in kwargs['backbone']: | |
| scale_factor = 2 | |
| else: | |
| scale_factor = 1 | |
| self.predict_head_out = nn.Sequential( | |
| nn.Conv2d(in_channels=last_channels, out_channels=last_channels * scale_factor, kernel_size=3, stride=1, padding=1), | |
| nn.BatchNorm2d(last_channels * scale_factor), | |
| nn.LeakyReLU(), | |
| nn.Conv2d(in_channels=last_channels * scale_factor, out_channels=n_classes, kernel_size=3, stride=1, padding=1), | |
| ) | |
| self.loss_att_branch = nn.Sequential( | |
| nn.Conv2d(in_channels=self.reduce_dim_out * 2, out_channels=64, kernel_size=3, stride=1, padding=1), | |
| nn.BatchNorm2d(64), | |
| nn.LeakyReLU(), | |
| nn.Conv2d(in_channels=64, out_channels=n_classes, kernel_size=3, stride=1, padding=1), | |
| ) | |
| def forward(self, x, *args, **kwargs): | |
| x, endpoints = self.res_backbone(x) | |
| # reduce channel 512 to 128 | |
| x_reduced_channel = self.reduce_channel_b5(x) # B 128 h w | |
| prob_s_map = self.get_s_probability(x_reduced_channel) | |
| prob_c_map = self.get_c_probability(x_reduced_channel) # B C 1 1 | |
| x_att_s = x_reduced_channel * prob_s_map | |
| x_att_c = x_reduced_channel * prob_c_map | |
| output_cat = [] | |
| if 'X' in self.model_pattern: | |
| output_cat.append(x_reduced_channel) | |
| if 'A' in self.model_pattern: | |
| output_cat.append(x_att_s) | |
| output_cat.append(x_att_c) | |
| if 'S' in self.model_pattern: | |
| # spatial pos embedding | |
| prob_s_vector = rearrange(prob_s_map, 'b c h w -> b (h w) c') | |
| x_vec_s = rearrange(x_reduced_channel, 'b c h w -> b (h w) c') | |
| # get top k, k = 16 * 16 // 4 x_b5_reduced_channel_vector | |
| _, indices_s = torch.topk(prob_s_vector, k=self.top_k_s, dim=1, sorted=False) # B K 1 | |
| indices_s = repeat(indices_s, 'b k m -> b k (m c)', c=self.reduce_dim_out) | |
| x_s_vec_topk = torch.gather(x_vec_s, 1, indices_s) # B K 128 | |
| if self.encoder_pos or self.decoder_pos: | |
| s_pos_embedding = self.spatial_embedding_h + self.spatial_embedding_w # 1 128 16 16 | |
| s_pos_embedding = repeat(s_pos_embedding, 'm c h w -> (b m) c h w', b=x.size(0)) | |
| s_pos_embedding_vec = rearrange(s_pos_embedding, 'b c h w -> b (h w) c') | |
| s_pos_embedding_vec_topk = torch.gather(s_pos_embedding_vec, 1, indices_s) # B K 128 | |
| if self.encoder_pos is True: | |
| pos_encoder = s_pos_embedding_vec_topk | |
| else: | |
| pos_encoder = None | |
| # b5 encoder and decoder op | |
| tf_encoder_s_x = self.tf_encoder_spatial_b5( | |
| q_s=x_s_vec_topk, k_s=None, v_s=None, pos_emb=pos_encoder | |
| ) | |
| if self.decoder_pos is True: | |
| pos_decoder = s_pos_embedding_vec_topk | |
| else: | |
| pos_decoder = None | |
| tf_decoder_s_x = self.tf_decoder_spatial_b5( | |
| q_s=x_vec_s, k_s=tf_encoder_s_x, v_s=None, | |
| pos_emb=pos_decoder | |
| ) # B (16*16) 128 | |
| # B 128 16 16 | |
| tf_decoder_s_x = rearrange(tf_decoder_s_x, 'b (h w) c -> b c h w', h=self.f_map_size) | |
| output_cat.append(tf_decoder_s_x) | |
| if 'C' in self.model_pattern: | |
| # channel attention ops | |
| prob_c_vec = rearrange(prob_c_map, 'b c h w -> b c (h w)') | |
| x_vec_c = rearrange(x_reduced_channel, 'b c h w -> b c (h w)') | |
| # get top k, k = 128 // 4 = 32 | |
| _, indices_c = torch.topk(prob_c_vec, k=self.top_k_c, dim=1, sorted=True) # b k 1 | |
| indices_c = repeat(indices_c, 'b k m -> b k (m c)', c=self.f_map_size ** 2) | |
| x_vec_c_topk = torch.gather(x_vec_c, 1, indices_c) # B K 256 | |
| if self.encoder_pos or self.decoder_pos: | |
| c_pos_embedding_vec = repeat(self.channel_embedding, 'm len c -> (m b) len c', b=x.size(0)) | |
| c_pos_embedding_vec_topk = torch.gather(c_pos_embedding_vec, 1, indices_c) # B K 256 | |
| if self.encoder_pos is True: | |
| pos_encoder = c_pos_embedding_vec_topk | |
| else: | |
| pos_encoder = None | |
| # b5 encoder and decoder op | |
| tf_encoder_c_x = self.tf_encoder_channel_b5( | |
| q_s=x_vec_c_topk, k_s=None, v_s=None, | |
| pos_emb=pos_encoder | |
| ) | |
| if self.decoder_pos is True: | |
| pos_decoder = c_pos_embedding_vec_topk | |
| else: | |
| pos_decoder = None | |
| tf_decoder_c_x = self.tf_decoder_channel_b5( | |
| q_s=x_vec_c, k_s=tf_encoder_c_x, v_s=None, | |
| pos_emb=pos_decoder | |
| ) # B 128 (16*16) | |
| # B 128 16 16 | |
| tf_decoder_c_x = rearrange(tf_decoder_c_x, 'b c (h w) -> b c h w', h=self.f_map_size) | |
| output_cat.append(tf_decoder_c_x) | |
| x_cat = torch.cat(output_cat, dim=1) | |
| x_cat = self.before_predict_head_conv(x_cat) | |
| x = self.double_conv1(self.pixel_shuffle1(x_cat)) | |
| x = self.double_conv2(self.pixel_shuffle2(x)) | |
| logits = self.predict_head_out(x) | |
| att_output = torch.cat([x_att_s, x_att_c], dim=1) | |
| att_branch_output = self.loss_att_branch(att_output) | |
| return logits, att_branch_output | |