Spaces:
Build error
Build error
| # Two types of reconstructionl layers: 1. original residual layers, 2. residual layers with contrast and adaptive attention(CCA layer) | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.init as init | |
| def initialize_weights(net_l, scale=1): | |
| if not isinstance(net_l, list): | |
| net_l = [net_l] | |
| for net in net_l: | |
| for m in net.modules(): | |
| if isinstance(m, nn.Conv2d): | |
| init.kaiming_normal_(m.weight, a=0, mode='fan_in') | |
| m.weight.data *= scale # for residual block | |
| if m.bias is not None: | |
| m.bias.data.zero_() | |
| elif isinstance(m, nn.Linear): | |
| init.kaiming_normal_(m.weight, a=0, mode='fan_in') | |
| m.weight.data *= scale | |
| if m.bias is not None: | |
| m.bias.data.zero_() | |
| elif isinstance(m, nn.BatchNorm2d): | |
| init.constant_(m.weight, 1) | |
| init.constant_(m.bias.data, 0.0) | |
| def make_layer(block, n_layers): | |
| layers = [] | |
| for _ in range(n_layers): | |
| layers.append(block()) | |
| return nn.Sequential(*layers) | |
| class ResidualBlock_noBN(nn.Module): | |
| """Residual block w/o BN | |
| ---Conv-ReLU-Conv-+- | |
| |________________| | |
| """ | |
| def __init__(self, nf=64): | |
| super(ResidualBlock_noBN, self).__init__() | |
| self.conv1 = nn.Conv2d(nf, nf, kernel_size=3, stride=1, padding=1, bias=True) | |
| self.conv2 = nn.Conv2d(nf, nf, kernel_size=3, stride=1, padding=1, bias=True) | |
| self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True) | |
| def forward(self, x): | |
| """ | |
| Args: | |
| x: with shape of [b, c, t, h, w] | |
| Returns: processed features with shape [b, c, t, h, w] | |
| """ | |
| identity = x | |
| out = self.lrelu(self.conv1(x)) | |
| out = self.conv2(out) | |
| out = identity + out | |
| # Remove ReLU at the end of the residual block | |
| # http://torch.ch/blog/2016/02/04/resnets.html | |
| return out | |
| class ResBlock_noBN_new(nn.Module): | |
| def __init__(self, nf): | |
| super(ResBlock_noBN_new, self).__init__() | |
| self.c1 = nn.Conv3d(nf, nf // 4, kernel_size=(1, 3, 3), stride=1, padding=(0, 1, 1), bias=True) | |
| self.d1 = nn.Conv3d(nf // 4, nf // 4, kernel_size=(1, 3, 3), stride=1, padding=(0, 1, 1), | |
| bias=True) # dilation rate=1 | |
| self.d2 = nn.Conv3d(nf // 4, nf // 4, kernel_size=(1, 3, 3), stride=1, padding=(0, 2, 2), dilation=(1, 2, 2), | |
| bias=True) # dilation rate=2 | |
| self.d3 = nn.Conv3d(nf // 4, nf // 4, kernel_size=(1, 3, 3), stride=1, padding=(0, 4, 4), dilation=(1, 4, 4), | |
| bias=True) # dilation rate=4 | |
| self.d4 = nn.Conv3d(nf // 4, nf // 4, kernel_size=(1, 3, 3), stride=1, padding=(0, 8, 8), dilation=(1, 8, 8), | |
| bias=True) # dilation rate=8 | |
| self.act = nn.LeakyReLU(negative_slope=0.2, inplace=True) | |
| self.c2 = nn.Conv3d(nf, nf, kernel_size=(1, 3, 3), stride=1, padding=(0, 1, 1), bias=True) | |
| def forward(self, x): | |
| output1 = self.act(self.c1(x)) | |
| d1 = self.d1(output1) | |
| d2 = self.d2(output1) | |
| d3 = self.d3(output1) | |
| d4 = self.d4(output1) | |
| add1 = d1 + d2 | |
| add2 = add1 + d3 | |
| add3 = add2 + d4 | |
| combine = torch.cat([d1, add1, add2, add3], dim=1) | |
| output2 = self.c2(self.act(combine)) | |
| output = x + output2 | |
| # remove ReLU at the end of the residual block | |
| # http://torch.ch/blog/2016/02/04/resnets.html | |
| return output | |
| class CCALayer(nn.Module): #############################################3 new | |
| '''Residual block w/o BN | |
| --conv--contrast-conv--x--- | |
| | \--mean--| | | |
| |___________________| | |
| ''' | |
| def __init__(self, nf=64): | |
| super(CCALayer, self).__init__() | |
| self.conv1 = nn.Conv2d(nf, nf, 3, 1, 1, bias=True) | |
| self.conv2 = nn.Conv2d(nf, nf, 3, 1, 1, bias=True) | |
| self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True) | |
| self.conv_du = nn.Sequential( | |
| nn.Conv2d(nf, 4, 1, padding=0, bias=True), | |
| nn.ReLU(inplace=True), | |
| nn.Conv2d(4, nf, 1, padding=0, bias=True), | |
| nn.Tanh() # change from `Sigmoid` to `Tanh` to make the output between -1 and 1 | |
| ) | |
| self.contrast = stdv_channels | |
| self.avg_pool = nn.AdaptiveAvgPool2d(1) | |
| # initialization | |
| initialize_weights([self.conv1, self.conv_du], 0.1) | |
| def forward(self, x): | |
| identity = x | |
| out = self.lrelu(self.conv1(x)) | |
| out = self.conv2(out) | |
| out = self.contrast(out) + self.avg_pool(out) | |
| out_channel = self.conv_du(out) | |
| out_channel = out_channel * out | |
| out_last = out_channel + identity | |
| return out_last | |
| def mean_channels(F): | |
| assert (F.dim() == 4), 'Your dim is {} bit not 4'.format(F.dim()) | |
| spatial_sum = F.sum(3, keepdim=True).sum(2, keepdim=True) | |
| return spatial_sum / (F.size(2) * F.size(3)) # 对每一个channel都求其特征图的高和宽的平均值 | |
| def stdv_channels(F): | |
| assert F.dim() == 4, 'Your dim is {} bit not 4'.format(F.dim()) | |
| F_mean = mean_channels(F) | |
| F_variance = (F - F_mean).pow(2).sum(3, keepdim=True).sum(2, keepdim=True) / (F.size(2) * F.size(3)) | |
| return F_variance.pow(0.5) | |