def __init__(self, dim, context_dim, hid_dim,
mask=0, realify=nn_.softplus):
super(BlockAffineFlow, self).__init__()
self.mask = mask
self.dim = dim
self.realify = realify
self.gpu = True
self.hid = nn_.WNBilinear(dim, context_dim, hid_dim)
self.mean = nn_.ResLinear(hid_dim, dim)
self.lstd = nn_.ResLinear(hid_dim, dim)
def __init__(self, dimc, act=nn.ELU()):
super(MNISTConvEnc, self).__init__()
self.enc = nn.Sequential(
nn_.ResConv2d(1, 16, 3, 2, padding=1, activation=act), act,
nn_.ResConv2d(16, 16, 3, 1, padding=1, activation=act), act,
nn_.ResConv2d(16, 32, 3, 2, padding=1, activation=act), act,
nn_.ResConv2d(32, 32, 3, 1, padding=1, activation=act), act,
nn_.ResConv2d(32, 32, 3, 2, padding=1, activation=act), act,
nn_.Reshape((-1, 32 * 4 * 4)), nn_.ResLinear(32 * 4 * 4, dimc),
act)
def __init__(self, dimz, dimc, act=nn.ELU()):
super(MNISTConvDec, self).__init__()
self.dec = nn.Sequential(
nn_.ResLinear(dimz, dimc),
act,
nn_.ResLinear(dimc, 32 * 4 * 4),
act,
nn_.Reshape((-1, 32, 4, 4)),
nn.Upsample(scale_factor=2, mode='bilinear'),
nn_.ResConv2d(32, 32, 3, 1, padding=1, activation=act),
act,
nn_.ResConv2d(32, 32, 3, 1, padding=1, activation=act),
act,
nn_.slicer[:, :, :-1, :-1],
nn.Upsample(scale_factor=2, mode='bilinear'),
nn_.ResConv2d(32, 16, 3, 1, padding=1, activation=act),
act,
nn_.ResConv2d(16, 16, 3, 1, padding=1, activation=act),
act,
nn.Upsample(scale_factor=2, mode='bilinear'),
nn_.ResConv2d(16, 1, 3, 1, padding=1, activation=act),
)