networks/w_encoder_arch.py

import torch
import torch.nn as nn
import torch.nn.functional as F
import math
from .prior_arch import PixelNorm, EqualLinear


class BasicBlock(nn.Module):
    def __init__(self, inplanes, planes, stride=1, downsample=None):
        super(BasicBlock, self).__init__()
        self.conv1 = conv1x1(inplanes, planes)
        self.gn1 = GroupNorm(planes)
        self.relu = nn.LeakyReLU(0.2, inplace=True)
        self.conv2 = conv3x3(planes, planes, stride)
        self.gn2 = GroupNorm(planes)
        self.downsample = downsample
        self.stride = stride

    def forward(self, x):
        residual = x
        out = self.conv1(x)
        out = self.gn1(out)
        out = self.relu(out)
        out = self.conv2(out)
        out = self.gn2(out)
        if self.downsample is not None:
            residual = self.downsample(x)
        out += residual
        out = self.relu(out)
        return out
        
class WEncoder(nn.Module):
    def __init__(self, block=BasicBlock, layers=[3, 4, 6, 6, 3], strides=[2,1,2,1,2]):
        self.inplanes = 32
        super(WEncoder, self).__init__()
        self.conv1 = nn.Conv2d(3, 32, kernel_size=3, stride=1, padding=1,
                               bias=False)
        self.relu = nn.LeakyReLU(0.2, inplace=True)

        feature_out_dim = 512
        self.layer1 = self._make_layer(block, 32, layers[0], stride=strides[0])
        self.layer2 = self._make_layer(block, 64, layers[1], stride=strides[1])
        self.layer3 = self._make_layer(block, 128, layers[2], stride=strides[2])
        self.layer4 = self._make_layer(block, 256, layers[3], stride=strides[3])
        self.layer5 = self._make_layer(block, feature_out_dim, layers[4], stride=strides[4])


        self.down_h = 1
        for stride in strides:
            self.down_h *= stride
        self.size_h = 32 // self.down_h

    
        self.feature2w = nn.Sequential(
            PixelNorm(),
            EqualLinear(self.size_h*self.size_h*feature_out_dim, 512, bias=True, bias_init_val=0, lr_mul=1,
                    activation='fused_lrelu'),
            EqualLinear(512, 512, bias=True, bias_init_val=0, lr_mul=1,
                    activation='fused_lrelu')
            # EqualLinear(self.size_h*self.size_h*feature_out_dim, 512, bias=True),
            # EqualLinear(512, 512, bias=True)
        )

        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                m.weight.data.normal_(0, math.sqrt(2. / n))


    def _make_layer(self, block, planes, blocks, stride=1):
        downsample = None
        if stride != 1 or self.inplanes != planes:
            downsample = nn.Sequential(
                nn.Conv2d(self.inplanes, planes,
                          kernel_size=1, stride=stride, bias=False),
            )
                # GroupNorm(planes),

        layers = []
        layers.append(block(self.inplanes, planes, stride, downsample))
        self.inplanes = planes
        for i in range(1, blocks):
            layers.append(block(self.inplanes, planes))

        return nn.Sequential(*layers)

    def _check_outliers(self, crop_feature, target_width):
        _, _, H, W = crop_feature.size()
        if W != target_width:
            return F.interpolate(crop_feature, size=(H, target_width), mode='bilinear', align_corners=True)
        else:
            return crop_feature


    def forward(self, x, locs):
        # lr = x.clone()
        x = self.conv1(x)
        x = self.relu(x)
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)
        x = self.layer5(x) # B, 512, 4, 64, 17M parameters
  
        B, C, H, W = x.size()

        # lr = F.interpolate(lr, (x.size(2), x.size(3)))
        w_b = []
        for b in range(locs.size(0)): #locs: 0~2048
            w_c = []
            for c in range(locs.size(1)):
                if locs[b][c] < 2048:
                    center_loc = (locs[b][c]/4/self.down_h).int() # from 32*512 to 4*64
                    start_x = max(0, center_loc-self.size_h//2)
                    end_x = min(center_loc+self.size_h//2, 512//self.down_h)
                    
                    # crop_feature = x[b:b+1, :, :, start_x:end_x].clone()
                    # crop_feature = self._check_outliers(crop_feature, self.size_h) # 1, 512, 4, 4 or 1, 512, 8, 8
                    
                    if end_x - start_x != self.size_h:
                        bgfill = torch.zeros((B, C, H, self.size_h), dtype=x.dtype, layout=x.layout, device=x.device)
                        bgfill[:, :, :, self.size_h//2 - (center_loc - start_x):self.size_h//2 - (center_loc - start_x) + end_x - start_x] += x[b:b+1, :, :, start_x:end_x].clone() 
                        crop_feature = bgfill.clone()
                    else:
                        crop_feature = x[b:b+1, :, :, start_x:end_x].clone()
                    w = self.feature2w(crop_feature.view(1, -1)) # 1*512
                    w_c.append(w.squeeze(0))

                else:
                    w_c.append(w.squeeze(0).detach()*0)

            w_c = torch.stack(w_c, dim=0)
            w_b.append(w_c)
        w_b = torch.stack(w_b, dim=0)
        
        return w_b #, lr



def GroupNorm(in_channels):
    return torch.nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=False)


def conv1x1(in_planes, out_planes, stride=1):
    return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)


def conv3x3(in_planes, out_planes, stride=1):
    "3x3 convolution with padding"
    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
                     padding=1, bias=False)
     
    


def _upsample_add(x, y):
    '''Upsample and add two feature maps.
    Args:
        x: (Variable) top feature map to be upsampled.
        y: (Variable) lateral feature map.
    Returns:
        (Variable) added feature map.
    Note in PyTorch, when input size is odd, the upsampled feature map
    with `F.upsample(..., scale_factor=2, mode='nearest')`
    maybe not equal to the lateral feature map size.
    e.g.
    original input size: [N,_,15,15] ->
    conv2d feature map size: [N,_,8,8] ->
    upsampled feature map size: [N,_,16,16]
    So we choose bilinear upsample which supports arbitrary output sizes.
    '''
    _, _, H, W = y.size()
    return F.interpolate(x, size=(H, W), mode='bilinear', align_corners=True) + y




if __name__ == '__main__':
    from .helper_arch import network_param

    device = 'cuda'
    input = torch.randn(2, 3, 32, 512).to(device) #

    test_list = [64]
    for i in range(1, 8):
        test_list.append(64+128*i)
    
    for i in range(8, 16):
        test_list.append(2048)

    
    
    locs = torch.Tensor(test_list).unsqueeze(0)
    locs = locs.repeat(2, 1).to(device)
    net = WEncoder().to(device)
    '''
    strides=[2,1,2,1,1] output h is 8
    Encoder is 12.97M
    F2W+Encoder is 17.04 M

    strides=[2,1,2,1,2] output h is 4
    Encoder is 12.97M
    F2W is 4.46 M

    '''

    output = net(input, locs)
    print([input.size(), output.size(), locs.size(), network_param(net)])
    #[torch.Size([2, 3, 32, 512]), torch.Size([2, 16, 512]), torch.Size([2, 16]), 17.43344]

    # import numpy as np
    # import cv2
    # sr_results = lr[0].permute(1, 2, 0)
    # sr_results = sr_results.float().cpu().numpy()
    
    # cv2.imwrite('./tmp.png', sr_results)