def __init__(self, z_channels):
super(DenseNetEncoderCoreColorImage64x64, self).__init__(z_channels)
self.main = nn.Sequential(
DenseNet(self.nc, 15, 3),
# state [48, 64, 64]
Conv2dWeightNorm(48, 48, 3, 2, 1, bias=True),
nn.ELU(),
# state [48, 32, 32]
DenseNet(48, 16, 3),
# state [96, 32, 32]
Conv2dWeightNorm(96, 96, 3, 2, 1, bias=True),
nn.ELU(),
# state [96, 16, 16]
DenseNet(96, 16, 3),
# state [144, 16, 16]
Conv2dWeightNorm(144, 144, 3, 2, 1, bias=True),
nn.ELU(),
# state [144, 8, 8]
DenseNet(144, 16, 3),
# state [192, 8, 8]
Conv2dWeightNorm(192, 96, 1, 1, bias=True),
nn.ELU(),
# state [96, 8, 8]
Conv2dWeightNorm(96, 2 * self.z_channels, 1, 1, bias=True),
# [2 * z_channels, 8, 8]
)
def __init__(self, z_channels):
super(DenseNetEncoderCoreColorImage32x32, self).__init__()
self.z_channels = z_channels
self.H = 8
self.W = 8
self.nc = 3
self.main = nn.Sequential(
DenseNet(self.nc, 15, 3),
# state [48, 32, 32]
Conv2dWeightNorm(48, 48, 3, 2, 1, bias=True),
nn.ELU(),
# state [48, 16, 16]
DenseNet(48, 16, 3),
# state [96, 16, 16]
Conv2dWeightNorm(96, 96, 3, 2, 1, bias=True),
nn.ELU(),
# state [96, 8, 8]
DenseNet(96, 16, 6),
# state [192, 8, 8]
Conv2dWeightNorm(192, 96, 1, 1, bias=True),
nn.ELU(),
# state [96, 8, 8]
Conv2dWeightNorm(96, 2 * self.z_channels, 1, 1, bias=True),
# [2 * z_channels, 8, 8]
)
def __init__(self, nz, mode, ngpu=1):
super(PixelCNNDecoderBinaryImage28x28, self).__init__(nz, ngpu=ngpu)
self.nc = 1
self.fm_latent = 4
self.img_latent = 28 * 28 * self.fm_latent
self.z_transform = nn.Sequential(
LinearWeightNorm(nz, self.img_latent),
ReShape((self.fm_latent, 28, 28)),
nn.ELU(),
)
if mode == 'small':
kernal_sizes = [7, 7, 7, 5, 5, 3, 3]
elif mode == 'large':
kernal_sizes = [7, 7, 7, 7, 7, 5, 5, 5, 5, 3, 3, 3, 3]
else:
raise ValueError('unknown mode: %s' % mode)
hidden_channels = 64
self.core = nn.Sequential(
PixelCNN(self.nc + self.fm_latent, hidden_channels,
len(kernal_sizes), kernal_sizes, self.nc),
Conv2dWeightNorm(hidden_channels, hidden_channels, 1, bias=True),
nn.ELU(),
Conv2dWeightNorm(hidden_channels, self.nc, 1, bias=True),
nn.Sigmoid(),
)
if ngpu > 1:
self.z_transform = nn.DataParallel(self.z_transform,
device_ids=list(range(ngpu)))
self.core = nn.DataParallel(self.core,
device_ids=list(range(ngpu)))
def __init__(self, in_channels, kernel_size):
super(PixelCNNBlock, self).__init__()
self.mask_type = 'B'
padding = kernel_size // 2
out_channels = in_channels // 2
self.main = nn.Sequential(
Conv2dWeightNorm(in_channels, out_channels, 1, bias=True),
nn.ELU(),
MaskedConv2d(out_channels, out_channels, kernel_size, mask_type='B', padding=padding, bias=True),
nn.ELU(),
Conv2dWeightNorm(out_channels, in_channels, 1, bias=True),
)
self.activation = nn.ELU()
def __init__(self, inplanes, growth_rate):
super(DenseNetBlock, self).__init__()
self.main = nn.Sequential(
Conv2dWeightNorm(inplanes,
4 * growth_rate,
kernel_size=1,
bias=True), nn.ELU(),
conv3x3(4 * growth_rate, growth_rate))
def __init__(self, inplanes, planes, stride=1):
super(ResNetBlock, self).__init__()
self.conv1 = conv3x3(inplanes, planes, stride)
self.activation = nn.ELU()
self.conv2 = conv3x3(planes, planes)
downsample = None
if stride != 1 or inplanes != planes:
downsample = Conv2dWeightNorm(inplanes, planes,
kernel_size=1, stride=stride, bias=True)
self.downsample = downsample
self.stride = stride
def __init__(self, in_channels, h_channels=0, dropout=0.0, activation=concat_elu):
super(GatedResnetBlock, self).__init__()
assert activation in [elu, concat_elu]
gain = 1 if activation == elu else 2
self.activation = activation
self.down_conv1 = DownShiftConv2d(gain * in_channels, in_channels, kernel_size=(2, 3), bias=True)
self.down_conv2 = DownShiftConv2d(gain * in_channels, in_channels, kernel_size=(2, 3), bias=True)
self.down_conv3 = DownShiftConv2d(gain * in_channels, 2 * in_channels, kernel_size=(2, 3), bias=True)
self.down_right_conv1 = DownRightShiftConv2d(gain * in_channels, in_channels, kernel_size=(2, 2), bias=True)
self.down_right_conv2 = DownRightShiftConv2d(gain * in_channels, in_channels, kernel_size=(2, 2), bias=True)
self.nin = Conv2dWeightNorm(gain * in_channels, in_channels, kernel_size=(1, 1))
self.down_right_conv3 = DownRightShiftConv2d(gain * in_channels, 2 * in_channels, kernel_size=(2, 2), bias=True)
if h_channels:
self.h_conv1 = Conv2dWeightNorm(gain * h_channels, in_channels, kernel_size=(3, 3), padding=1)
self.h_conv2 = Conv2dWeightNorm(gain * in_channels, 2 * in_channels, kernel_size=(3, 3), padding=1)
else:
self.h_conv1 = None
self.h_conv2 = None
self.dropout = nn.Dropout(p=dropout)
def __init__(self, nc, z_channels, h_channels, hidden_channels, levels, num_resnets, nmix, dropout=0., activation='concat_elu'):
super(_PixelCNNPPCore, self).__init__()
in_channels = z_channels
blocks = []
for i in range(levels - 1):
blocks.append(('level%d' % i, ConvTranspose2dWeightNorm(in_channels, in_channels // 2, 3, 2, 1, 1, bias=True)))
blocks.append(('elu%d' %i, nn.ELU()))
in_channels = in_channels // 2
blocks.append(('h_level', Conv2dWeightNorm(in_channels, h_channels, 1)))
self.z_transform = nn.Sequential(OrderedDict(blocks))
self.core = PixelCNNPP(levels, nc, hidden_channels, num_resnets, h_channels, dropout=dropout, activation=activation)
self.output = nn.Sequential(
nn.ELU(),
# state [hidden_channels, H, W]
Conv2dWeightNorm(hidden_channels, hidden_channels, 1, bias=True),
nn.ELU(),
# state [hidden_channels * 2, H, W]
Conv2dWeightNorm(hidden_channels, (nc * 3 + 1) * nmix, 1, bias=True)
# state [10 * nmix, H, W]
)
def conv3x3(in_planes, out_planes, stride=1):
"3x3 convolution with padding"
return Conv2dWeightNorm(in_planes, out_planes,
kernel_size=3, stride=stride,
padding=1, bias=True)
def __init__(self, num_filters):
super(DownSampling, self).__init__()
self.conv = Conv2dWeightNorm(num_filters, num_filters * 2, kernel_size=(3, 3), stride=(2, 2), padding=1)
def __init__(self, num_filters, activation=concat_elu):
super(NINBlock, self).__init__()
assert activation in [elu, concat_elu]
gain = 1 if activation == elu else 2
self.activation = activation
self.nin = Conv2dWeightNorm(gain * num_filters, num_filters, kernel_size=(1, 1))