|
import torch |
|
import torch.nn as nn |
|
import torch.nn.functional as F |
|
|
|
import cliport.utils.utils as utils |
|
|
|
|
|
class IdentityBlock(nn.Module): |
|
def __init__(self, in_planes, filters, kernel_size, stride=1, final_relu=True, batchnorm=True): |
|
super(IdentityBlock, self).__init__() |
|
self.final_relu = final_relu |
|
self.batchnorm = batchnorm |
|
|
|
filters1, filters2, filters3 = filters |
|
self.conv1 = nn.Conv2d(in_planes, filters1, kernel_size=1, bias=False) |
|
self.bn1 = nn.BatchNorm2d(filters1) if self.batchnorm else nn.Identity() |
|
self.conv2 = nn.Conv2d(filters1, filters2, kernel_size=kernel_size, dilation=1, |
|
stride=stride, padding=1, bias=False) |
|
self.bn2 = nn.BatchNorm2d(filters2) if self.batchnorm else nn.Identity() |
|
self.conv3 = nn.Conv2d(filters2, filters3, kernel_size=1, bias=False) |
|
self.bn3 = nn.BatchNorm2d(filters3) if self.batchnorm else nn.Identity() |
|
|
|
def forward(self, x): |
|
out = F.relu(self.bn1(self.conv1(x))) |
|
out = F.relu(self.bn2(self.conv2(out))) |
|
out = self.bn3(self.conv3(out)) |
|
out += x |
|
if self.final_relu: |
|
out = F.relu(out) |
|
return out |
|
|
|
|
|
class ConvBlock(nn.Module): |
|
def __init__(self, in_planes, filters, kernel_size, stride=1, final_relu=True, batchnorm=True): |
|
super(ConvBlock, self).__init__() |
|
self.final_relu = final_relu |
|
self.batchnorm = batchnorm |
|
|
|
filters1, filters2, filters3 = filters |
|
self.conv1 = nn.Conv2d(in_planes, filters1, kernel_size=1, bias=False) |
|
self.bn1 = nn.BatchNorm2d(filters1) if self.batchnorm else nn.Identity() |
|
self.conv2 = nn.Conv2d(filters1, filters2, kernel_size=kernel_size, dilation=1, |
|
stride=stride, padding=1, bias=False) |
|
self.bn2 = nn.BatchNorm2d(filters2) if self.batchnorm else nn.Identity() |
|
self.conv3 = nn.Conv2d(filters2, filters3, kernel_size=1, bias=False) |
|
self.bn3 = nn.BatchNorm2d(filters3) if self.batchnorm else nn.Identity() |
|
|
|
self.shortcut = nn.Sequential( |
|
nn.Conv2d(in_planes, filters3, |
|
kernel_size=1, stride=stride, bias=False), |
|
nn.BatchNorm2d(filters3) if self.batchnorm else nn.Identity() |
|
) |
|
|
|
def forward(self, x): |
|
out = F.relu(self.bn1(self.conv1(x))) |
|
out = F.relu(self.bn2(self.conv2(out))) |
|
out = self.bn3(self.conv3(out)) |
|
out += self.shortcut(x) |
|
if self.final_relu: |
|
out = F.relu(out) |
|
return out |
|
|
|
|
|
class ResNet43_8s(nn.Module): |
|
def __init__(self, input_shape, output_dim, cfg, device, preprocess): |
|
super(ResNet43_8s, self).__init__() |
|
self.input_shape = input_shape |
|
self.input_dim = input_shape[-1] |
|
self.output_dim = output_dim |
|
self.cfg = cfg |
|
self.device = device |
|
self.batchnorm = self.cfg['train']['batchnorm'] |
|
self.preprocess = preprocess |
|
|
|
self.layers = self._make_layers() |
|
|
|
def _make_layers(self): |
|
layers = nn.Sequential( |
|
|
|
nn.Conv2d(self.input_dim, 64, stride=1, kernel_size=3, padding=1), |
|
nn.BatchNorm2d(64) if self.batchnorm else nn.Identity(), |
|
nn.ReLU(True), |
|
|
|
|
|
ConvBlock(64, [64, 64, 64], kernel_size=3, stride=1, batchnorm=self.batchnorm), |
|
IdentityBlock(64, [64, 64, 64], kernel_size=3, stride=1, batchnorm=self.batchnorm), |
|
|
|
ConvBlock(64, [128, 128, 128], kernel_size=3, stride=2, batchnorm=self.batchnorm), |
|
IdentityBlock(128, [128, 128, 128], kernel_size=3, stride=1, batchnorm=self.batchnorm), |
|
|
|
ConvBlock(128, [256, 256, 256], kernel_size=3, stride=2, batchnorm=self.batchnorm), |
|
IdentityBlock(256, [256, 256, 256], kernel_size=3, stride=1, batchnorm=self.batchnorm), |
|
|
|
ConvBlock(256, [512, 512, 512], kernel_size=3, stride=2, batchnorm=self.batchnorm), |
|
IdentityBlock(512, [512, 512, 512], kernel_size=3, stride=1, batchnorm=self.batchnorm), |
|
|
|
|
|
ConvBlock(512, [256, 256, 256], kernel_size=3, stride=1, batchnorm=self.batchnorm), |
|
IdentityBlock(256, [256, 256, 256], kernel_size=3, stride=1, batchnorm=self.batchnorm), |
|
nn.UpsamplingBilinear2d(scale_factor=2), |
|
|
|
ConvBlock(256, [128, 128, 128], kernel_size=3, stride=1, batchnorm=self.batchnorm), |
|
IdentityBlock(128, [128, 128, 128], kernel_size=3, stride=1, batchnorm=self.batchnorm), |
|
nn.UpsamplingBilinear2d(scale_factor=2), |
|
|
|
ConvBlock(128, [64, 64, 64], kernel_size=3, stride=1, batchnorm=self.batchnorm), |
|
IdentityBlock(64, [64, 64, 64], kernel_size=3, stride=1, batchnorm=self.batchnorm), |
|
nn.UpsamplingBilinear2d(scale_factor=2), |
|
|
|
|
|
ConvBlock(64, [16, 16, self.output_dim], kernel_size=3, stride=1, |
|
final_relu=False, batchnorm=self.batchnorm), |
|
IdentityBlock(self.output_dim, [16, 16, self.output_dim], kernel_size=3, stride=1, |
|
final_relu=False, batchnorm=self.batchnorm), |
|
) |
|
return layers |
|
|
|
def forward(self, x): |
|
x = self.preprocess(x, dist='transporter') |
|
|
|
out = self.layers(x) |
|
return out |
