def __init__(self):
super(InferenceCore, self).__init__()
self.downsample_x = nn.Conv2d(1,
1,
kernel_size=4,
stride=4,
padding=0,
bias=False)
self.upsample_v = nn.ConvTranspose2d(8,
8,
kernel_size=16,
stride=16,
padding=0,
bias=False)
self.upsample_r = nn.ConvTranspose2d(256,
256,
kernel_size=16,
stride=16,
padding=0,
bias=False)
self.downsample_u = nn.Conv2d(128,
128,
kernel_size=4,
stride=4,
padding=0,
bias=False)
self.core = Conv2dLSTMCell(1 + 8 + 256 + 2 * 128,
128,
kernel_size=5,
stride=1,
padding=2)
def __init__(self):
super(GenerationCore, self).__init__()
self.upsample_v = nn.ConvTranspose2d(8,
8,
kernel_size=16,
stride=16,
padding=0,
bias=False)
self.upsample_r = nn.ConvTranspose2d(256,
256,
kernel_size=16,
stride=16,
padding=0,
bias=False)
self.core = Conv2dLSTMCell(8 + 256 + 1,
128,
kernel_size=5,
stride=1,
padding=2)
self.upsample_h = nn.ConvTranspose2d(128,
128,
kernel_size=4,
stride=4,
padding=0,
bias=False)
def __init__(self):
super(InferenceCore, self).__init__()
self.downsample_x = nn.Conv2d(3,
3,
kernel_size=4,
stride=4,
padding=0,
bias=False)
self.upsample_v = nn.ConvTranspose2d(7,
7,
kernel_size=8,
stride=8,
padding=0,
bias=False)
self.upsample_r = nn.ConvTranspose2d(64,
64,
kernel_size=16,
stride=16,
padding=0,
bias=False)
self.downsample_u = nn.Conv2d(64,
64,
kernel_size=4,
stride=4,
padding=0,
bias=False)
self.core = Conv2dLSTMCell(3 + 7 + 64 + 2 * 64,
64,
kernel_size=5,
stride=1,
padding=2)
def __init__(self):
super(GenerationCore, self).__init__()
self.upsample_v = nn.ConvTranspose2d(7,
7,
kernel_size=8,
stride=8,
padding=0,
bias=False)
self.upsample_r = nn.ConvTranspose2d(64,
64,
kernel_size=16,
stride=16,
padding=0,
bias=False)
self.core = Conv2dLSTMCell(7 + 64 + 3,
64,
kernel_size=5,
stride=1,
padding=2)
self.upsample_h = nn.ConvTranspose2d(64,
64,
kernel_size=4,
stride=4,
padding=0,
bias=False)
def __init__(self, nf_to_hidden, stride_to_obs, nf_to_obs, nf_dec, nf_z,
nf_v):
super(Renderer, self).__init__()
self.conv = nn.Conv2d(nf_to_obs,
nf_dec,
kernel_size=stride_to_obs,
stride=stride_to_obs)
self.lstm = Conv2dLSTMCell(nf_z + nf_v + nf_dec,
nf_to_hidden,
kernel_size=5,
stride=1,
padding=2)
self.transconv = nn.ConvTranspose2d(nf_to_hidden,
nf_to_obs,
kernel_size=stride_to_obs,
stride=stride_to_obs)
def __init__(self, stride_to_hidden, nf_to_hidden, nf_enc, nf_z):
super(Prior, self).__init__()
self.conv1 = nn.Conv2d(32,
nf_enc,
kernel_size=stride_to_hidden,
stride=stride_to_hidden)
self.lstm = Conv2dLSTMCell(nf_enc + nf_z,
nf_to_hidden,
kernel_size=5,
stride=1,
padding=2)
self.conv2 = nn.Conv2d(nf_to_hidden,
2 * nf_z,
kernel_size=5,
stride=1,
padding=2)