def __init__(self, in_planes: int, stride: int = 1):
super(TemporalResidualBlock, 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)
def __init__(self,
in_planes: int,
out_planes: int,
stride: int = 1,
downsample: t.Optional[nn.Module] = None):
super(SpatialResidualBlock, self).__init__()
# Both self.conv1 and self.downsample layers downsample the input when stride != 1
self.conv1 = tc.conv3x3(in_planes, out_planes, stride)
self.bn1 = nn.BatchNorm2d(out_planes)
self.relu = nn.ReLU(inplace=True)
self.conv2 = tc.conv3x3(out_planes, out_planes)
self.bn2 = nn.BatchNorm2d(out_planes)
self.downsample = downsample
self.stride = stride
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,
in_planes: int,
out_planes: int,
stride: int = 1,
upsample: t.Optional[nn.Module] = None):
super(TransSpatialResidualBlock, self).__init__()
self.conv1 = tc.conv3x3(in_planes, in_planes)
self.bn1 = nn.BatchNorm2d(in_planes)
self.relu = nn.ReLU(inplace=True)
self.upsample = upsample
self.conv2 = tc.conv3x3(in_planes, out_planes, stride)
if upsample is not None and stride != 1:
self.conv2 = tc.t_conv3x3(in_planes,
out_planes,
stride=stride,
padding=1)
else:
self.conv2 = tc.conv3x3(in_planes, out_planes, stride)
self.bn2 = nn.BatchNorm2d(out_planes)