| from networks import * | |
| class PosADANet(nn.Module): | |
| def encode(self, shp): | |
| device = self.omega.device | |
| B, _, H, W = shp | |
| row = torch.arange(H).to(device) / H | |
| enc_row1 = torch.sin(self.omega[None, :] * row[:, None]) | |
| enc_row2 = torch.cos(self.omega[None, :] * row[:, None]) | |
| rows = torch.cat([enc_row1.unsqueeze(1).repeat((1, W, 1)), enc_row2.unsqueeze(1).repeat((1, W, 1))], dim=-1) | |
| col = torch.arange(W).to(device) / W | |
| enc_col1 = torch.sin(self.omega[None, :] * col[:, None]) | |
| enc_col2 = torch.cos(self.omega[None, :] * col[:, None]) | |
| cols = torch.cat([enc_col1.unsqueeze(0).repeat((H, 1, 1)), enc_col2.unsqueeze(0).repeat((H, 1, 1))], dim=-1) | |
| encoding = torch.cat([rows, cols], dim=-1) | |
| encoding = encoding.permute(2, 0, 1).unsqueeze(0).repeat((B, 1, 1, 1)) | |
| return encoding | |
| def get_encoding(self, x): | |
| shp1 = x.shape | |
| singelton = self.positional_encoding is not None\ | |
| and self.positional_encoding.shape[0] == shp1[0] and self.positional_encoding.shape[2:] == shp1[2:] | |
| if singelton: | |
| return self.positional_encoding | |
| self.positional_encoding = self.encode(x.shape) | |
| return self.positional_encoding | |
| def __init__(self, input_channels, output_channels, n_style, bilinear=True, padding='zero', full_ada=True, | |
| nfreq=20, magnitude=10): | |
| super(PosADANet, self).__init__() | |
| factor = 2 if bilinear else 1 | |
| self.omega = nn.Parameter(torch.rand(nfreq) * magnitude) | |
| self.omega.requires_grad = False | |
| self.positional_encoding = None | |
| self.full_ada = full_ada | |
| self.style_encoder = FullyConnected(n_style, W_SIZE, layers=6) | |
| self.padding = padding | |
| self.input_channels = input_channels + nfreq * 4 | |
| self.n_classes = output_channels | |
| self.bilinear = bilinear | |
| self.channels = [512 // factor, 256 // factor, 128 // factor] | |
| self.inc = DoubleConv(self.input_channels, 64) | |
| self.down1 = Down(64, 128, padding=padding, ada=self.full_ada) | |
| self.down2 = Down(128, 256, padding=padding, ada=self.full_ada) | |
| self.down3 = Down(256, 512, padding=padding, ada=self.full_ada) | |
| self.down4 = Down(512, 1024 // factor, padding=padding, ada=self.full_ada) | |
| self.up1 = Up(1024, 512 // factor, bilinear, ada=True, padding=padding) | |
| self.up2 = Up(512, 256 // factor, bilinear, ada=True, padding=padding) | |
| self.up3 = Up(256, 128 // factor, bilinear, ada=True, padding=padding) | |
| self.up4 = Up(128, 64, bilinear, padding=padding, ada=True) | |
| self.outc = OutConv(64, output_channels, padding=padding) | |
| def forward(self, x, style): | |
| w = self.style_encoder(style) | |
| encoding = self.get_encoding(x) | |
| x = torch.cat([x, encoding], dim=1) | |
| x1 = self.inc(x) | |
| if self.full_ada: | |
| x2 = self.down1(x1, w=w) | |
| x3 = self.down2(x2, w=w) | |
| x4 = self.down3(x3, w=w) | |
| x5 = self.down4(x4, w=w) | |
| else: | |
| x2 = self.down1(x1) | |
| x3 = self.down2(x2) | |
| x4 = self.down3(x3) | |
| x5 = self.down4(x4) | |
| x = self.up1(x5, x4, w=w) | |
| x = self.up2(x, x3, w=w) | |
| x = self.up3(x, x2, w=w) | |
| x = self.up4(x, x1, w=w) | |
| logits = self.outc(x) | |
| return logits | |