""" modules.py - This file stores the rather boring network blocks. x - usually means features that only depends on the image g - usually means features that also depends on the mask. They might have an extra "group" or "num_objects" dimension, hence batch_size * num_objects * num_channels * H * W The trailing number of a variable usually denote the stride """ import torch import torch.nn as nn import torch.nn.functional as F from model.group_modules import * from model import resnet from model.cbam import CBAM class FeatureFusionBlock(nn.Module): def __init__(self, x_in_dim, g_in_dim, g_mid_dim, g_out_dim): super().__init__() self.distributor = MainToGroupDistributor() self.block1 = GroupResBlock(x_in_dim + g_in_dim, g_mid_dim) self.attention = CBAM(g_mid_dim) self.block2 = GroupResBlock(g_mid_dim, g_out_dim) def forward(self, x, g): batch_size, num_objects = g.shape[:2] g = self.distributor(x, g) g = self.block1(g) r = self.attention(g.flatten(start_dim=0, end_dim=1)) r = r.view(batch_size, num_objects, *r.shape[1:]) g = self.block2(g + r) return g class HiddenUpdater(nn.Module): # Used in the decoder, multi-scale feature + GRU def __init__(self, g_dims, mid_dim, hidden_dim): super().__init__() self.hidden_dim = hidden_dim self.g16_conv = GConv2D(g_dims[0], mid_dim, kernel_size=1) self.g8_conv = GConv2D(g_dims[1], mid_dim, kernel_size=1) self.g4_conv = GConv2D(g_dims[2], mid_dim, kernel_size=1) self.transform = GConv2D( mid_dim + hidden_dim, hidden_dim * 3, kernel_size=3, padding=1 ) nn.init.xavier_normal_(self.transform.weight) def forward(self, g, h): g = ( self.g16_conv(g[0]) + self.g8_conv(downsample_groups(g[1], ratio=1 / 2)) + self.g4_conv(downsample_groups(g[2], ratio=1 / 4)) ) g = torch.cat([g, h], 2) # defined slightly differently than standard GRU, # namely the new value is generated before the forget gate. # might provide better gradient but frankly it was initially just an # implementation error that I never bothered fixing values = self.transform(g) forget_gate = torch.sigmoid(values[:, :, : self.hidden_dim]) update_gate = torch.sigmoid(values[:, :, self.hidden_dim : self.hidden_dim * 2]) new_value = torch.tanh(values[:, :, self.hidden_dim * 2 :]) new_h = forget_gate * h * (1 - update_gate) + update_gate * new_value return new_h class HiddenReinforcer(nn.Module): # Used in the value encoder, a single GRU def __init__(self, g_dim, hidden_dim): super().__init__() self.hidden_dim = hidden_dim self.transform = GConv2D( g_dim + hidden_dim, hidden_dim * 3, kernel_size=3, padding=1 ) nn.init.xavier_normal_(self.transform.weight) def forward(self, g, h): g = torch.cat([g, h], 2) # defined slightly differently than standard GRU, # namely the new value is generated before the forget gate. # might provide better gradient but frankly it was initially just an # implementation error that I never bothered fixing values = self.transform(g) forget_gate = torch.sigmoid(values[:, :, : self.hidden_dim]) update_gate = torch.sigmoid(values[:, :, self.hidden_dim : self.hidden_dim * 2]) new_value = torch.tanh(values[:, :, self.hidden_dim * 2 :]) new_h = forget_gate * h * (1 - update_gate) + update_gate * new_value return new_h class ValueEncoder(nn.Module): def __init__(self, value_dim, hidden_dim, single_object=False): super().__init__() self.single_object = single_object network = resnet.resnet18(pretrained=True, extra_dim=1 if single_object else 2) self.conv1 = network.conv1 self.bn1 = network.bn1 self.relu = network.relu # 1/2, 64 self.maxpool = network.maxpool self.layer1 = network.layer1 # 1/4, 64 self.layer2 = network.layer2 # 1/8, 128 self.layer3 = network.layer3 # 1/16, 256 self.distributor = MainToGroupDistributor() self.fuser = FeatureFusionBlock(1024, 256, value_dim, value_dim) if hidden_dim > 0: self.hidden_reinforce = HiddenReinforcer(value_dim, hidden_dim) else: self.hidden_reinforce = None def forward(self, image, image_feat_f16, h, masks, others, is_deep_update=True): # image_feat_f16 is the feature from the key encoder if not self.single_object: g = torch.stack([masks, others], 2) else: g = masks.unsqueeze(2) g = self.distributor(image, g) batch_size, num_objects = g.shape[:2] g = g.flatten(start_dim=0, end_dim=1) g = self.conv1(g) g = self.bn1(g) # 1/2, 64 g = self.maxpool(g) # 1/4, 64 g = self.relu(g) g = self.layer1(g) # 1/4 g = self.layer2(g) # 1/8 g = self.layer3(g) # 1/16 g = g.view(batch_size, num_objects, *g.shape[1:]) g = self.fuser(image_feat_f16, g) if is_deep_update and self.hidden_reinforce is not None: h = self.hidden_reinforce(g, h) return g, h class KeyEncoder(nn.Module): def __init__(self): super().__init__() network = resnet.resnet50(pretrained=True) self.conv1 = network.conv1 self.bn1 = network.bn1 self.relu = network.relu # 1/2, 64 self.maxpool = network.maxpool self.res2 = network.layer1 # 1/4, 256 self.layer2 = network.layer2 # 1/8, 512 self.layer3 = network.layer3 # 1/16, 1024 def forward(self, f): x = self.conv1(f) x = self.bn1(x) x = self.relu(x) # 1/2, 64 x = self.maxpool(x) # 1/4, 64 f4 = self.res2(x) # 1/4, 256 f8 = self.layer2(f4) # 1/8, 512 f16 = self.layer3(f8) # 1/16, 1024 return f16, f8, f4 class UpsampleBlock(nn.Module): def __init__(self, skip_dim, g_up_dim, g_out_dim, scale_factor=2): super().__init__() self.skip_conv = nn.Conv2d(skip_dim, g_up_dim, kernel_size=3, padding=1) self.distributor = MainToGroupDistributor(method="add") self.out_conv = GroupResBlock(g_up_dim, g_out_dim) self.scale_factor = scale_factor def forward(self, skip_f, up_g): skip_f = self.skip_conv(skip_f) g = upsample_groups(up_g, ratio=self.scale_factor) g = self.distributor(skip_f, g) g = self.out_conv(g) return g class KeyProjection(nn.Module): def __init__(self, in_dim, keydim): super().__init__() self.key_proj = nn.Conv2d(in_dim, keydim, kernel_size=3, padding=1) # shrinkage self.d_proj = nn.Conv2d(in_dim, 1, kernel_size=3, padding=1) # selection self.e_proj = nn.Conv2d(in_dim, keydim, kernel_size=3, padding=1) nn.init.orthogonal_(self.key_proj.weight.data) nn.init.zeros_(self.key_proj.bias.data) def forward(self, x, need_s, need_e): shrinkage = self.d_proj(x) ** 2 + 1 if (need_s) else None selection = torch.sigmoid(self.e_proj(x)) if (need_e) else None return self.key_proj(x), shrinkage, selection class Decoder(nn.Module): def __init__(self, val_dim, hidden_dim): super().__init__() self.fuser = FeatureFusionBlock(1024, val_dim + hidden_dim, 512, 512) if hidden_dim > 0: self.hidden_update = HiddenUpdater([512, 256, 256 + 1], 256, hidden_dim) else: self.hidden_update = None self.up_16_8 = UpsampleBlock(512, 512, 256) # 1/16 -> 1/8 self.up_8_4 = UpsampleBlock(256, 256, 256) # 1/8 -> 1/4 self.pred = nn.Conv2d(256, 1, kernel_size=3, padding=1, stride=1) def forward(self, f16, f8, f4, hidden_state, memory_readout, h_out=True): batch_size, num_objects = memory_readout.shape[:2] if self.hidden_update is not None: g16 = self.fuser(f16, torch.cat([memory_readout, hidden_state], 2)) else: g16 = self.fuser(f16, memory_readout) g8 = self.up_16_8(f8, g16) g4 = self.up_8_4(f4, g8) logits = self.pred(F.relu(g4.flatten(start_dim=0, end_dim=1))) if h_out and self.hidden_update is not None: g4 = torch.cat( [g4, logits.view(batch_size, num_objects, 1, *logits.shape[-2:])], 2 ) hidden_state = self.hidden_update([g16, g8, g4], hidden_state) else: hidden_state = None logits = F.interpolate( logits, scale_factor=4, mode="bilinear", align_corners=False ) logits = logits.view(batch_size, num_objects, *logits.shape[-2:]) return hidden_state, logits