Listening to the Pixels - VAN

van.py

import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision
from functools import partial

class DilatedResnet18MaxPool(nn.Module):
    def __init__(self, fc_dim=64, conv_size=3):
        super(DilatedResnet18MaxPool, self).__init__()

        orig_resnet = torchvision.models.resnet18(pretrained=True)
        orig_resnet.layer4.apply(
                partial(self._nostride_dilate, dilate=2))

        self.features = nn.Sequential(
            *list(orig_resnet.children())[:-2])

        self.fc = nn.Conv2d(
            512, fc_dim, kernel_size=conv_size, padding=conv_size//2)

    def _nostride_dilate(self, m, dilate):
        classname = m.__class__.__name__
        if classname.find('Conv') != -1:
            # the convolution with stride
            if m.stride == (2, 2):
                m.stride = (1, 1)
                if m.kernel_size == (3, 3):
                    m.dilation = (dilate//2, dilate//2)
                    m.padding = (dilate//2, dilate//2)
            # other convoluions
            else:
                if m.kernel_size == (3, 3):
                    m.dilation = (dilate, dilate)
                    m.padding = (dilate, dilate)

    def forward(self, x, pool=True):
        x = self.features(x)
        x = self.fc(x)

        if not pool:
            return x

        x = F.adaptive_max_pool2d(x, 1)
        x = x.view(x.size(0), x.size(1))
        return x

    def forward_multi(self, x, pool=True):
        (B, C, T, H, W) = x.size()
        x = x.permute(0, 2, 1, 3, 4).contiguous()
        x = x.view(B*T, C, H, W)

        x = self.features(x)
        x = self.fc(x)

        (_, C, H, W) = x.size()
        x = x.view(B, T, C, H, W)
        x = x.permute(0, 2, 1, 3, 4)

        if not pool:
            return x

        x = F.adaptive_max_pool3d(x, 1)
        x = x.view(B, C)
        return x