def __init__(self, in_planes: int, stride: int = 1):
super(VarTemporalResidualBlock, self).__init__()
self.in_planes = in_planes
self.stride = stride
self.conv1 = tc.conv3x3(in_planes, in_planes, stride)
self.bn1 = nn.BatchNorm2d(in_planes)
self.relu = nn.ReLU(inplace=True)
self.conv2 = tc.conv3x3(in_planes, in_planes)
self.bn2 = nn.BatchNorm2d(in_planes)
self.mean = tc.conv1x1(in_planes, in_planes)
self.var = nn.Sequential(tc.conv1x1(in_planes, in_planes),
nn.Softplus())
self.rsample = models.common.ReparameterizedSample()
def _make_transpose(self,
out_planes: int,
blocks: int,
stride: int = 1) -> nn.Module:
if stride != 1:
upsample = nn.Sequential(
nn.Upsample(scale_factor=2, mode='bicubic',
align_corners=True),
tc.conv3x3(self.in_planes,
out_planes * TransSpatialResidualBlock.expansion,
1),
nn.BatchNorm2d(out_planes),
)
elif self.in_planes != out_planes:
upsample = nn.Sequential(
tc.conv1x1(self.in_planes,
out_planes * TransSpatialResidualBlock.expansion,
stride),
nn.BatchNorm2d(out_planes),
)
else:
upsample = None
layers = []
for _ in range(0, blocks - 1):
layers.append(
TransSpatialResidualBlock(self.in_planes, self.in_planes))
layers.append(
TransSpatialResidualBlock(self.in_planes, out_planes, stride,
upsample))
self.in_planes = out_planes * TransSpatialResidualBlock.expansion
if self.flow:
self.in_planes *= 2
return nn.Sequential(*layers)
def __init__(self, out_planes: t.Tuple[int, ...], bottleneck_planes: int):
super(SpatialResNetEncoder, self).__init__()
self.in_planes = out_planes[0]
self.out_planes = out_planes
self.bottleneck_planes = bottleneck_planes
self.conv1 = nn.Conv2d(3,
self.out_planes[0],
kernel_size=7,
stride=2,
padding=3,
bias=False)
self.bn1 = nn.BatchNorm2d(self.out_planes[0])
self.relu = nn.ReLU(inplace=True)
self.max_pool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
self.layers = nn.ModuleList()
self.layers.append(self._make_layer(self.out_planes[1], 2))
for out_plane in self.out_planes[2:]:
self.layers.append(self._make_layer(out_plane, 2, stride=2))
self.bottleneck = nn.Sequential(
tc.conv1x1(self.in_planes, self.bottleneck_planes),
nn.BatchNorm2d(self.bottleneck_planes),
nn.ReLU(inplace=True),
)
def _make_layer(self,
planes: int,
blocks: int,
stride: int = 1) -> nn.Module:
downsample = None
if stride != 1 or self.in_planes != planes * SpatialResidualBlock.expansion:
downsample = nn.Sequential(
tc.conv1x1(self.in_planes,
planes * SpatialResidualBlock.expansion, stride),
nn.BatchNorm2d(planes * SpatialResidualBlock.expansion),
)
layers = [
SpatialResidualBlock(self.in_planes, planes, stride, downsample)
]
self.in_planes = planes * SpatialResidualBlock.expansion
for _ in range(1, blocks):
layers.append(SpatialResidualBlock(self.in_planes, planes, 1))
return nn.Sequential(*layers)
def __init__(self, in_planes: int, out_planes: t.Tuple[int, ...],
flow: bool):
super(SpatialResNetDecoder, self).__init__()
self.in_planes = in_planes
self.out_planes = out_planes
self.flow = flow
self.debottleneck = nn.Sequential(
tc.conv1x1(self.in_planes, self.out_planes[0]),
nn.BatchNorm2d(self.out_planes[0]),
nn.ReLU(inplace=True),
)
self.in_planes = self.out_planes[0]
if self.flow:
self.in_planes *= 2
self.layers = nn.ModuleList()
for out_plane in self.out_planes[1:-1]:
self.layers.append(self._make_transpose(out_plane, 2, stride=2))
self.layers.append(
self._make_transpose(self.out_planes[-1], 2, stride=1))
if self.flow:
self.final = nn.ConvTranspose2d(self.out_planes[-1],
2,
kernel_size=3,
stride=1,
padding=1,
bias=True)
else:
self.final = nn.ConvTranspose2d(self.out_planes[-1],
3,
kernel_size=3,
stride=1,
padding=1,
bias=True)
self.final_act = nn.Sigmoid()