| import math |
|
|
| import torch |
| import torch.nn as nn |
|
|
|
|
| def cylindrical_conv_pad(x, w_pad): |
| return torch.cat([x[..., -w_pad:], x, x[..., :w_pad]], axis=-1) |
|
|
|
|
| class CylindricalConv2D(nn.Conv2d): |
| """ |
| UNet with cylinderical boundary conditions |
| """ |
|
|
| def __init__( |
| self, |
| in_channels: int, |
| out_channels: int, |
| kernel_size: int, |
| stride: int, |
| ): |
| super().__init__( |
| in_channels=in_channels, |
| out_channels=out_channels, |
| kernel_size=kernel_size, |
| stride=stride, |
| ) |
|
|
| assert self.kernel_size[0] % 2 == 1 |
| assert self.kernel_size[1] % 2 == 1 |
|
|
| self.h_pad = self.kernel_size[0] // 2 |
| self.w_pad = self.kernel_size[1] // 2 |
|
|
| def forward(self, x: torch.Tensor): |
| x = nn.functional.pad(x, (0, 0, self.h_pad, self.h_pad)) |
| return super().forward(cylindrical_conv_pad(x, self.w_pad)) |
|
|
|
|
| class CylindricalConvTranspose2D(nn.ConvTranspose2d): |
| def __init__( |
| self, |
| in_channels: int, |
| out_channels: int, |
| kernel_size: int, |
| stride: int, |
| ): |
| super().__init__( |
| in_channels=in_channels, |
| out_channels=out_channels, |
| kernel_size=kernel_size, |
| stride=stride, |
| ) |
|
|
| assert self.kernel_size[0] % 2 == 1 |
| assert self.kernel_size[1] % 2 == 1 |
|
|
| self.sh, self.sw = self.stride |
| self.kh, self.kw = self.kernel_size |
|
|
| self.h_pad = math.ceil(((self.sh - 1) + 2 * (self.kh // 2)) / self.sh) |
| self.w_pad = math.ceil(((self.sw - 1) + 2 * (self.kw // 2)) / self.sw) |
|
|
| self.h0 = self.sh * self.h_pad - (self.sh - 1) + (self.kh // 2) |
| self.w0 = self.sw * self.w_pad - (self.sw - 1) + (self.kw // 2) |
|
|
| self._bias = nn.Parameter(10**-3 * torch.randn(out_channels)) |
|
|
| def forward(self, x: torch.Tensor): |
|
|
| Nh = x.shape[2] * self.sh |
| Nw = x.shape[3] * self.sw |
|
|
| x = cylindrical_conv_pad(x, self.w_pad) |
| x = nn.functional.pad(x, (0, 0, self.h_pad, self.h_pad)) |
|
|
| x = super().forward(x) |
|
|
| return x[:, :, self.h0 : self.h0 + Nh, self.w0 : self.w0 + Nw] |
|
|
|
|
| class Down(nn.Module): |
|
|
| def __init__( |
| self, |
| in_channels, |
| out_channels, |
| p=0, |
| film=False, |
| down=True, |
| attn=False, |
| ): |
|
|
| super().__init__() |
|
|
| self.film = film |
| self.attn = attn |
|
|
| self.conv_1 = CylindricalConv2D( |
| in_channels, out_channels, kernel_size=3, stride=1 |
| ) |
| if down: |
| self.conv_2 = CylindricalConv2D( |
| out_channels, out_channels, kernel_size=3, stride=2 |
| ) |
| else: |
| self.conv_2 = CylindricalConv2D( |
| out_channels, out_channels, kernel_size=3, stride=1 |
| ) |
|
|
| self.bn_1 = nn.BatchNorm2d(out_channels) |
| self.bn_2 = nn.BatchNorm2d(out_channels) |
|
|
| self.activation = nn.GELU() |
|
|
| if film: |
|
|
| self.gamma_1 = torch.nn.Parameter( |
| torch.ones(10, out_channels, 1, 1), |
| ) |
| self.gamma_2 = torch.nn.Parameter( |
| torch.ones(10, out_channels, 1, 1), |
| ) |
| self.beta_1 = torch.nn.Parameter( |
| torch.zeros(10, out_channels, 1, 1), |
| ) |
| self.beta_2 = torch.nn.Parameter( |
| torch.zeros(10, out_channels, 1, 1), |
| ) |
|
|
| if self.attn: |
| self.mha = AttentionBlock(n_channels=out_channels, n_heads=8) |
|
|
| def forward(self, xi, film_index=None): |
|
|
| film_index = film_index[:, 0].int() |
|
|
| x = self.conv_1(xi) |
| x = self.bn_1(x) |
| if self.film: |
| g1 = torch.index_select(self.gamma_1, 0, film_index) |
| b1 = torch.index_select(self.beta_1, 0, film_index) |
|
|
| x = g1 * x + b1 |
|
|
| x = self.activation(x) |
| x = self.conv_2(x) |
| x = self.bn_2(x) |
| if self.film: |
| g2 = torch.index_select(self.gamma_2, 0, film_index) |
| b2 = torch.index_select(self.beta_2, 0, film_index) |
| x = g2 * x + b2 |
| x = self.activation(x) |
|
|
| if self.attn: |
| x = self.mha(x) |
|
|
| return x |
|
|
|
|
| class Up(nn.Module): |
|
|
| def __init__( |
| self, |
| in_channels, |
| out_channels, |
| p, |
| bilinear=False, |
| film=False, |
| stride=2, |
| attn=False, |
| ): |
| super().__init__() |
|
|
| self.film = film |
|
|
| if bilinear: |
| self.up = nn.Upsample( |
| scale_factor=2, |
| mode="bilinear", |
| align_corners=True, |
| ) |
| else: |
| self.up = CylindricalConvTranspose2D( |
| in_channels, out_channels, kernel_size=3, stride=stride |
| ) |
|
|
| self.conv = Down( |
| out_channels, |
| out_channels, |
| p=0, |
| film=film, |
| down=False, |
| attn=attn, |
| ) |
|
|
| def forward(self, x1, x2, film_index=None): |
| x1 = self.up(x1) |
| x1 = self.conv(x1, film_index=film_index) |
|
|
| if x1.shape[-1] != x2.shape[-1]: |
| x1 = x1[..., :, :-1] |
| if x1.shape[-2] != x2.shape[-2]: |
| x1 = x1[..., :-1, :] |
|
|
| return torch.cat([x2, x1], dim=1) |
|
|
|
|
| class Unet(nn.Module): |
|
|
| def __init__( |
| self, |
| in_channels, |
| out_channels, |
| div_factor=1, |
| p=0.0, |
| context=True, |
| film=False, |
| film_base=True, |
| ): |
| super(Unet, self).__init__() |
|
|
| self.n_channels = in_channels |
| self.bilinear = True |
| self.fp = nn.Softplus() |
| self.variances = nn.Parameter(torch.zeros([out_channels])) |
| self.context = context |
| self.film = film |
|
|
| m = 1 |
|
|
| self.down1 = Down( |
| self.n_channels, |
| m * 128 // div_factor, |
| p=0, |
| film=film, |
| attn=False, |
| ) |
| self.down2 = Down( |
| m * 128 // div_factor, |
| m * 256 // div_factor, |
| p=0, |
| film=film, |
| attn=False, |
| ) |
| self.down3 = Down( |
| m * 256 // div_factor, |
| m * 512 // div_factor, |
| p=0, |
| film=film, |
| attn=False, |
| ) |
| self.down4 = Down( |
| m * 512 // div_factor, |
| m * 512 // div_factor, |
| p=0, |
| film=film, |
| attn=False, |
| ) |
| self.up1 = Up( |
| m * 512 // div_factor, |
| m * 512 // div_factor, |
| p=0, |
| film=film, |
| attn=False, |
| ) |
| self.up2 = Up( |
| m * 1024 // div_factor, |
| m * 256 // div_factor, |
| p=0, |
| film=film, |
| attn=False, |
| ) |
| self.up3 = Up( |
| m * 512 // div_factor, |
| m * 128 // div_factor, |
| p=0, |
| film=film, |
| attn=False, |
| ) |
| self.up4 = Up( |
| m * 256 // div_factor, |
| m * 64 // div_factor, |
| p=0, |
| film=film, |
| attn=False, |
| ) |
|
|
| self.out = nn.Conv2d( |
| m * 64 // div_factor + in_channels, |
| out_channels, |
| kernel_size=1, |
| bias=False, |
| ) |
|
|
| def forward(self, x, film_index=None): |
|
|
| x1 = x.contiguous() |
| x2 = self.down1(x1, film_index=film_index) |
| x3 = self.down2(x2, film_index=film_index) |
| x4 = self.down3(x3, film_index=film_index) |
| x5 = self.down4(x4, film_index=film_index) |
| x = self.up1(x5, x4, film_index=film_index) |
| x = self.up2(x, x3, film_index=film_index) |
| x = self.up3(x, x2, film_index=film_index) |
| x = self.up4(x, x1, film_index=film_index) |
| x = self.out(x) |
|
|
| return x.permute(0, 2, 3, 1) |
|
|
