def __init__(self, dim_in, dim_out, spec_norm=False):
super(ResidualBlockUp, self).__init__()
if spec_norm:
self.stream1 = nn.Sequential(
nn.UpsamplingBilinear2d(scale_factor=2),
SpectralNorm(nn.Conv2d(dim_in, dim_out, kernel_size=1, stride=1, bias=False))
)
self.stream2 = nn.Sequential(
nn.BatchNorm2d(dim_in, affine=True, track_running_stats=True),
nn.ReLU(),
nn.UpsamplingBilinear2d(scale_factor=2),
SpectralNorm(nn.Conv2d(dim_in, dim_in, kernel_size=3, stride=1, padding=1, bias=False)),
nn.BatchNorm2d(dim_in, affine=True, track_running_stats=True),
nn.ReLU(),
SpectralNorm(nn.Conv2d(dim_in, dim_out, kernel_size=3, stride=1, padding=1, bias=False)),
)
else:
self.stream1 = nn.Sequential(
nn.UpsamplingBilinear2d(scale_factor=2),
nn.Conv2d(dim_in, dim_out, kernel_size=1, stride=1, bias=False)
)
self.stream2 = nn.Sequential(
nn.BatchNorm2d(dim_in, affine=True, track_running_stats=True),
nn.ReLU(),
nn.UpsamplingBilinear2d(scale_factor=2),
nn.Conv2d(dim_in, dim_in, kernel_size=3, stride=1, padding=1, bias=False),
nn.BatchNorm2d(dim_in, affine=True, track_running_stats=True),
nn.ReLU(),
nn.Conv2d(dim_in, dim_out, kernel_size=3, stride=1, padding=1, bias=False),
)
def __init__(self, dim_in, dim_out, spec_norm=False):
super(ConvDownBlock, self).__init__()
if spec_norm:
self.main = nn.Sequential(
nn.Conv2d(dim_in, dim_out, kernel_size=3, stride=1, padding=1, bias=False),
nn.BatchNorm2d(dim_out, affine=True, track_running_stats=True),
nn.LeakyReLU(0.01, inplace=False),
nn.Conv2d(dim_out, dim_out, kernel_size=3, stride=1, padding=1, bias=False),
nn.BatchNorm2d(dim_out, affine=True, track_running_stats=True),
nn.LeakyReLU(0.01, inplace=False),
nn.Conv2d(dim_out, dim_out, kernel_size=4, stride=2, padding=1, bias=False),
nn.BatchNorm2d(dim_out, affine=True, track_running_stats=True),
)
else:
self.main = nn.Sequential(
SpectralNorm(nn.Conv2d(dim_in, dim_out, kernel_size=3, stride=1, padding=1, bias=False)),
nn.BatchNorm2d(dim_out, affine=True, track_running_stats=True),
nn.LeakyReLU(0.01, inplace=False),
SpectralNorm(nn.Conv2d(dim_out, dim_out, kernel_size=3, stride=1, padding=1, bias=False)),
nn.BatchNorm2d(dim_out, affine=True, track_running_stats=True),
nn.LeakyReLU(0.01, inplace=False),
SpectralNorm(nn.Conv2d(dim_out, dim_out, kernel_size=4, stride=2, padding=1, bias=False)),
nn.BatchNorm2d(dim_out, affine=True, track_running_stats=True),
)
def __init__(self, batch_size=64, image_size=64, conv_dim=64):
super(Discriminator, self).__init__()
self.imsize = image_size
self.l1 = nn.Sequential(SpectralNorm(nn.Conv2d(3, conv_dim, 4, 2, 1)),
nn.LeakyReLU(0.1))
curr_dim = conv_dim
self.l2 = nn.Sequential(
SpectralNorm(nn.Conv2d(curr_dim, curr_dim * 2, 4, 2, 1)),
nn.LeakyReLU(0.1))
curr_dim *= 2
self.l3 = nn.Sequential(
SpectralNorm(nn.Conv2d(curr_dim, curr_dim * 2, 4, 2, 1)),
nn.LeakyReLU(0.1))
curr_dim *= 2
if self.imsize == 64:
self.l4 = nn.Sequential(
SpectralNorm(nn.Conv2d(curr_dim, curr_dim * 2, 4, 2, 1)),
nn.LeakyReLU(0.1))
curr_dim *= 2
self.last = nn.Sequential(nn.Conv2d(curr_dim, 1, 4))
self.attn1 = SelfAttention(256, 'relu')
self.attn2 = SelfAttention(512, 'relu')
def __init__(self, in_channels, out_channels, stride=1):
super(ResBlockDiscriminator, self).__init__()
self.conv1 = nn.Conv2d(in_channels, out_channels, 3, 1, padding=1)
self.conv2 = nn.Conv2d(out_channels, out_channels, 3, 1, padding=1)
nn.init.xavier_uniform(self.conv1.weight.data, 1.)
nn.init.xavier_uniform(self.conv2.weight.data, 1.)
if stride == 1:
self.model = nn.Sequential(nn.ReLU(), SpectralNorm(self.conv1),
nn.ReLU(), SpectralNorm(self.conv2))
else:
self.model = nn.Sequential(
nn.ReLU(), SpectralNorm(self.conv1), nn.ReLU(),
SpectralNorm(self.conv2),
nn.AvgPool2d(2, stride=stride, padding=0))
self.bypass = []
if in_channels != out_channels:
self.bypass_conv = nn.Conv2d(in_channels,
out_channels,
1,
1,
padding=0)
nn.init.xavier_uniform(self.bypass_conv.weight.data, 1.)
self.bypass.append(SpectralNorm(self.bypass_conv))
if stride != 1:
self.bypass.append(nn.AvgPool2d(2, stride=stride, padding=0))
# if in_channels == out_channels:
# self.bypass = nn.AvgPool2d(2, stride=stride, padding=0)
# else:
# self.bypass = nn.Sequential(
# SpectralNorm(nn.Conv2d(in_channels,out_channels, 1, 1, padding=0)),
# nn.AvgPool2d(2, stride=stride, padding=0)
# )
self.bypass = nn.Sequential(*self.bypass)
def __init__(self, in_channels, out_channels, stride=1):
super(ResBlockDiscriminator, self).__init__()
self.conv1 = nn.Conv2d(in_channels, out_channels, 3, 1, padding=1)
self.conv2 = nn.Conv2d(out_channels, out_channels, 3, 1, padding=1)
nn.init.xavier_uniform_(self.conv1.weight.data, 1.)
nn.init.xavier_uniform_(self.conv2.weight.data, 1.)
if stride == 1:
self.model = nn.Sequential(
nn.ReLU(),
SpectralNorm(self.conv1),
nn.ReLU(),
SpectralNorm(self.conv2)
)
else:
self.model = nn.Sequential(
nn.ReLU(),
SpectralNorm(self.conv1),
nn.ReLU(),
SpectralNorm(self.conv2),
nn.AvgPool2d(2, stride=stride, padding=0)
)
self.bypass = nn.Sequential()
if stride != 1:
self.bypass_conv = nn.Conv2d(in_channels,out_channels, 1, 1, padding=0)
nn.init.xavier_uniform_(self.bypass_conv.weight.data, np.sqrt(2))
self.bypass = nn.Sequential(
SpectralNorm(self.bypass_conv),
nn.AvgPool2d(2, stride=stride, padding=0)
)
def __init__(self, image_size=128, conv_dim=64, c_dim=5, repeat_num=6):
super(Discriminator, self).__init__()
layers = []
layers.append(
SpectralNorm(
nn.Conv2d(3, conv_dim, kernel_size=4, stride=2, padding=1)))
layers.append(nn.LeakyReLU(0.01, inplace=True))
curr_dim = conv_dim
for i in range(1, repeat_num):
layers.append(
SpectralNorm(
nn.Conv2d(curr_dim,
curr_dim * 2,
kernel_size=4,
stride=2,
padding=1)))
layers.append(nn.LeakyReLU(0.01, inplace=True))
curr_dim = curr_dim * 2
k_size = int(image_size / np.power(2, repeat_num))
self.main = nn.Sequential(*layers)
self.conv1 = nn.Conv2d(curr_dim,
1,
kernel_size=3,
stride=1,
padding=1,
bias=False)
self.conv2 = nn.Conv2d(curr_dim, c_dim, kernel_size=k_size, bias=False)
def __init__(self): super(Encoder, self).__init__() self.input_dim = 3 self.input_height = 64 self.input_width = 64 self.output_dim = 50 self.conv = nn.Sequential( SpectralNorm(nn.Conv2d(self.input_dim, 64, 3, 4, 2, bias=True)), #nn.BatchNorm2d(64), #nn.InstanceNorm2d(64, affine=True), nn.ReLU(), SpectralNorm(nn.Conv2d(64, 128, 4, 2, 1, bias=True)), #nn.BatchNorm2d(128), #nn.InstanceNorm2d(128, affine=True), nn.ReLU(), SpectralNorm(nn.Conv2d(128, 256, 4, 2, 1, bias=True)), #nn.BatchNorm2d(256), #nn.InstanceNorm2d(256, affine=True), nn.ReLU(), SpectralNorm(nn.Conv2d(256, 512, 4, 2, 1, bias=True)), #nn.BatchNorm2d(512), #nn.InstanceNorm2d(512, affine=True), nn.ReLU(), nn.Conv2d(512, self.output_dim, 4, 2, 1, bias=True), #nn.Sigmoid(), )
def __init__(self, num_cls): super(Discriminator, self).__init__() self.input_dim = 3 self.num_cls = num_cls self.conv = nn.Sequential( #64->32 SpectralNorm(nn.Conv2d(self.input_dim, 32, 4, 2, 1, bias=False)), #nn.BatchNorm2d(32), nn.LeakyReLU(), #32->16 SpectralNorm(nn.Conv2d(32, 64, 4, 2, 1, bias=False)), #nn.BatchNorm2d(64), nn.LeakyReLU(0.2), #16->8 SpectralNorm(nn.Conv2d(64, 128, 4, 2, 1, bias=False)), #nn.BatchNorm2d(128), nn.LeakyReLU(0.2), #8->4 SpectralNorm(nn.Conv2d(128, 256, 4, 2, 1, bias=False)), #nn.BatchNorm2d(256), nn.LeakyReLU(0.2), ) self.convCls = nn.Sequential( nn.Conv2d(256, self.num_cls, 4, bias=False) ) self.convGAN = nn.Sequential( nn.Conv2d(256, 1, 4, bias=False), nn.Sigmoid() )
def __init__(self, in_channels, out_channels, stride=1, activation="ELU"):
super(FirstResBlockDiscriminator, self).__init__()
self.conv1 = nn.Conv2d(in_channels, out_channels, 3, 1, padding=1)
self.conv2 = nn.Conv2d(out_channels, out_channels, 3, 1, padding=1)
self.bypass_conv = nn.Conv2d(in_channels, out_channels, 1, 1, padding=0)
nn.init.xavier_uniform_(self.conv1.weight.data, 1.)
nn.init.xavier_uniform_(self.conv2.weight.data, 1.)
nn.init.xavier_uniform_(self.bypass_conv.weight.data, np.sqrt(2))
if activation == "ELU":
Activation = nn.ELU
elif activation == "ReLU":
Activation = nn.ReLU
else:
NotImplementedError("now only implemented for ELU or ReLU")
# we don't want to apply ReLU activation to raw image before convolution transformation.
self.model = nn.Sequential(
SpectralNorm(self.conv1),
Activation(),
SpectralNorm(self.conv2),
nn.AvgPool2d(2)
)
self.bypass = nn.Sequential(
nn.AvgPool2d(2),
SpectralNorm(self.bypass_conv),
)
def __init__(self, batch_size, image_size=64, z_dim=100, conv_dim=64):
super(Generator, self).__init__()
self.imsize = image_size
repeat_num = int(np.log2(self.imsize)) - 3
mult = 2**repeat_num # 8
self.l1 = nn.Sequential(
SpectralNorm(nn.ConvTranspose2d(z_dim, conv_dim * mult, 4)),
nn.BatchNorm2d(conv_dim * mult), nn.ReLU())
curr_dim = conv_dim * mult
self.l2 = nn.Sequential(
SpectralNorm(nn.ConvTranspose2d(curr_dim, curr_dim // 2, 4, 2, 1)),
nn.BatchNorm2d(curr_dim // 2), nn.ReLU())
curr_dim //= 2
self.l3 = nn.Sequential(
SpectralNorm(nn.ConvTranspose2d(curr_dim, curr_dim // 2, 4, 2, 1)),
nn.BatchNorm2d(curr_dim // 2), nn.ReLU())
if self.imsize == 64:
curr_dim //= 2
self.l4 = nn.Sequential(
SpectralNorm(
nn.ConvTranspose2d(curr_dim, curr_dim // 2, 4, 2, 1)),
nn.BatchNorm2d(curr_dim // 2), nn.ReLU())
curr_dim //= 2
self.last = nn.Sequential(nn.ConvTranspose2d(curr_dim, 3, 4, 2, 1),
nn.Tanh())
self.attn1 = SelfAttention(128, 'relu')
self.attn2 = SelfAttention(64, 'relu')
def __init__(self, dim_in, dim_out, spec_norm=False):
super(ResidualBlockDown, self).__init__()
if spec_norm:
self.stream1 = nn.Sequential(
SpectralNorm(nn.Conv2d(dim_in, dim_out, kernel_size=1, stride=1, bias=False)),
nn.AvgPool2d(kernel_size=2, stride=2),
)
self.stream2 = nn.Sequential(
nn.ReLU(),
nn.Conv2d(dim_in, dim_in, kernel_size=3, stride=1, padding=1, bias=False),
nn.ReLU(),
SpectralNorm(nn.Conv2d(dim_in, dim_out, kernel_size=3, stride=1, padding=1, bias=False)),
nn.AvgPool2d(kernel_size=2, stride=2),
)
else:
self.stream1 = nn.Sequential(
nn.Conv2d(dim_in, dim_out, kernel_size=1, stride=1, bias=False),
nn.AvgPool2d(kernel_size=2, stride=2),
)
self.stream2 = nn.Sequential(
nn.ReLU(),
nn.Conv2d(dim_in, dim_in, kernel_size=3, stride=1, padding=1, bias=False),
nn.ReLU(),
nn.Conv2d(dim_in, dim_out, kernel_size=3, stride=1, padding=1, bias=False),
nn.AvgPool2d(kernel_size=2, stride=2),
)
def __init__(self, input_size, attention_size):
super(SelfAttentionPost, self).__init__()
self.attention_size = attention_size
self.gamma = nn.Parameter(torch.tensor(0.))
self.h = SpectralNorm(
nn.Conv2d(input_size, self.attention_size, 1, stride=1))
self.i = SpectralNorm(
nn.Conv2d(self.attention_size, input_size, 1, stride=1))
def forward(self, x):
if self.use_bn:
out = self.conv1(F.relu(self.bn1(x)))
out = self.conv2(F.relu(self.bn2(x)))
else:
out = SpectralNorm(self.conv1(F.relu(x)))
out = SpectralNorm(self.conv2(F.relu(x)))
return out + self.shortcut(x)
def __init__(self, input_size, attention_size):
super(SelfAttention, self).__init__()
self.attention_size = attention_size
#attention layers
self.f = SpectralNorm(
nn.Conv2d(input_size, attention_size, 1, stride=1))
self.g = SpectralNorm(
nn.Conv2d(input_size, attention_size, 1, stride=1))
self.input_size = input_size
def __init__(self, bit):
super(DisOrFuncf, self).__init__()
self.fc1 = SpectralNorm(nn.Linear(784, 512))
self.fc2 = SpectralNorm(nn.Linear(512, 256))
self.fc3 = SpectralNorm(nn.Linear(256, 1))
self.relu = nn.LeakyReLU(0.2)
self.sigm = nn.Sigmoid()
nc = 2**bit
self.a = self._build_mask(nc)
self.nc = nc
def __init__(self, latent_size):
super(Encoder, self).__init__()
self.latent_size = latent_size
# 80 x 80
self.conv1 = SpectralNorm(
nn.Conv2d(channels, 64, 3, stride=1, padding=(1, 1)))
# 80 x 80
self.conv2 = SpectralNorm(
nn.Conv2d(64, 64, 4, stride=2, padding=(1, 1)))
# 40 x 40
self.conv3 = SpectralNorm(
nn.Conv2d(64, 128, 3, stride=1, padding=(1, 1)))
# 40 x 40
self.conv4 = SpectralNorm(
nn.Conv2d(128, 128, 4, stride=2, padding=(1, 1)))
# 20 x 20
self.conv5 = SpectralNorm(
nn.Conv2d(128, 256, 3, stride=1, padding=(1, 1)))
# 20 x 20
self.conv6 = SpectralNorm(
nn.Conv2d(256, 256, 4, stride=2, padding=(1, 1)))
# 10 x 10
self.conv7 = SpectralNorm(
nn.Conv2d(256, 256, 4, stride=2, padding=(1, 1)))
# 5 x 5
self.conv8 = SpectralNorm(
nn.Conv2d(256, 512, 3, stride=1, padding=(0, 0)))
# 3 x 3
self.fc = SpectralNorm(nn.Linear(3 * 3 * 512, latent_size))
def __init__(self, input_dim=1, output_dim=1, input_size=32, class_num=10):
super(discriminator, self).__init__()
self.input_dim = input_dim
self.output_dim = output_dim
self.input_size = input_size
self.class_num = class_num
self.conv = nn.Sequential(
SpectralNorm(nn.Conv2d(3, 64, 3, 1, 1)),
nn.LeakyReLU(0.1),
SpectralNorm(nn.Conv2d(64, 64, 4, 2, 1)),
nn.LeakyReLU(0.1),
SpectralNorm(nn.Conv2d(64, 128, 3, 1, 1)),
nn.LeakyReLU(0.1),
SpectralNorm(nn.Conv2d(128, 128, 4, 2, 1)),
nn.LeakyReLU(0.1),
SpectralNorm(nn.Conv2d(128, 256, 3, 1, 1)),
nn.LeakyReLU(0.1),
SpectralNorm(nn.Conv2d(256, 256, 4, 2, 1)),
nn.LeakyReLU(0.1),
SpectralNorm(nn.Conv2d(256, 512, 3, 1, 1)),
)
self.dc = nn.Sequential(
SpectralNorm(nn.Linear(512 * 4 * 4, self.output_dim)),
nn.Sigmoid(),
)
self.cl = nn.Sequential(
SpectralNorm(nn.Linear(512 * 4 * 4, self.class_num)),
nn.Sigmoid(),
)
utils.initialize_weights(self)
def __init__(self):
super(Discriminator, self).__init__()
self.conv1 = SpectralNorm(
nn.Conv2d(channels, 64, 3, stride=1, padding=(2, 2)))
self.conv2 = SpectralNorm(
nn.Conv2d(64, 64, 4, stride=2, padding=(1, 1)))
self.conv3 = SpectralNorm(
nn.Conv2d(64, 128, 3, stride=1, padding=(1, 1)))
self.conv4 = SpectralNorm(
nn.Conv2d(128, 128, 4, stride=2, padding=(1, 1)))
self.conv5 = SpectralNorm(
nn.Conv2d(128, 256, 3, stride=1, padding=(1, 1)))
self.conv6 = SpectralNorm(
nn.Conv2d(256, 256, 4, stride=2, padding=(1, 1)))
self.attention_size = 32
self.att = SelfAttention(256, self.attention_size)
self.att_post = SelfAttentionPost(256, self.attention_size)
self.conv7 = SpectralNorm(
nn.Conv2d(256, 512, 3, stride=1, padding=(1, 1)))
self.embed = SpectralNorm(nn.Linear(num_classes, w_g * w_g * 512))
self.fc = SpectralNorm(nn.Linear(w_g * w_g * 512, 1))
def __init__(self):
super(Discriminator, self).__init__()
# 80 x 80
self.conv1 = SpectralNorm(
nn.Conv2d(channels, 64, 3, stride=1, padding=(1, 1)))
# 80 x 80
self.conv2 = SpectralNorm(
nn.Conv2d(64, 64, 4, stride=2, padding=(1, 1)))
# 40 x 40
self.conv3 = SpectralNorm(
nn.Conv2d(64, 128, 3, stride=1, padding=(1, 1)))
# 40 x 40
self.conv4 = SpectralNorm(
nn.Conv2d(128, 128, 4, stride=2, padding=(1, 1)))
# 20 x 20
self.conv5 = SpectralNorm(
nn.Conv2d(128, 256, 3, stride=1, padding=(1, 1)))
# 20 x 20
self.conv6 = SpectralNorm(
nn.Conv2d(256, 256, 4, stride=2, padding=(1, 1)))
# 10 x 10
self.conv7 = SpectralNorm(
nn.Conv2d(256, 256, 4, stride=2, padding=(1, 1)))
# 5 x 5
self.conv8 = SpectralNorm(
nn.Conv2d(256, 512, 3, stride=1, padding=(0, 0)))
# 3 x 3
self.fc = SpectralNorm(nn.Linear(w_g * w_g * 512, 1))
def __init__(self, size1, size2):
super(WDEstimator, self).__init__()
# Vanilla MLP
self.net = nn.Sequential(
SpectralNorm(nn.Linear(size1 + size2, 1024)),
SpectralNorm(nn.ReLU(True)),
SpectralNorm(nn.Linear(1024, 1024)),
SpectralNorm(nn.ReLU(True)),
SpectralNorm(nn.Linear(1024, 1)),
)
def __init__(self):
super().__init__()
# Bx1x64x64
self.conv1 = SpectralNorm(nn.Conv2d(NOISE_DIM, 32, (3, 3), stride=2, padding=0))
# Bx32x32x32
self.conv2 = SpectralNorm(nn.Conv2d(32, 32, (3, 3), stride=2, padding=0))
# Bx32x16x16
self.conv3 = nn.Conv2d(32, 32, (3,3), stride=2, padding=0)
self.fc1 = nn.Linear(32*7*7, 1)
self.cuda()
def __init__(self, latent_size, num_actions):
super().__init__()
# Input: State + Action
# Output: State
self.latent_size = latent_size
# Skip connections from output of 1 to input of 6, and output of 2 to input of 5
self.conv1 = SpectralNorm(nn.Conv2d(latent_size + num_actions, 128, (3,3), stride=1, padding=2, padding_mode='circular'))
self.conv2 = SpectralNorm(nn.Conv2d(128, 128, (3,3), stride=1, padding=2, padding_mode='circular'))
self.conv3 = SpectralNorm(nn.Conv2d(128, 128, (3,3), stride=1, padding=2, padding_mode='circular'))
self.conv4 = SpectralNorm(nn.Conv2d(128, 128, (3,3), stride=1, padding=2, padding_mode='circular'))
self.conv5 = SpectralNorm(nn.Conv2d(128 + 128, 128, (3,3), stride=1, padding=2, padding_mode='circular'))
self.conv6 = nn.Conv2d(128 + 128, latent_size, (3,3), stride=1, padding=2, padding_mode='circular')
self.cuda()
def __init__(self, latent_size, color_channels):
super().__init__()
self.latent_size = latent_size
self.color_channels = color_channels
# Bx1x64x64
self.conv1 = SpectralNorm(nn.Conv2d(color_channels * ENCODER_INPUT_FRAMES, 128, (3,3), stride=1, padding=1))
self.bn_conv1 = nn.BatchNorm2d(128)
# Bx8x32x32
self.conv2 = SpectralNorm(nn.Conv2d(128, 128, (3,3), stride=1, padding=1))
self.conv3 = SpectralNorm(nn.Conv2d(128, 128, (3,3), stride=1, padding=1))
self.conv4 = nn.Conv2d(128, latent_size, (3,3), stride=1, padding=1)
# Bxlatent_size
self.cuda()
def __init__(self):
super(Critic, self).__init__()
if MODE == 'sgan':
self.main = nn.Sequential(
SpectralNorm(nn.Linear(64, 256)), nn.LeakyReLU(0.1),
SpectralNorm(nn.Linear(256, 1, bias=False)), nn.Sigmoid())
elif MODE == 'wgan':
self.main = nn.Sequential(
SpectralNorm(nn.Linear(64, 256)), nn.LeakyReLU(0.1),
SpectralNorm(nn.Linear(256, 1, bias=False)))
# Print summary if VERBOSE is True
if VERBOSE:
self.summary()
def __init__(self, in_channels, out_channels, stride=1):
super(FirstResBlockDiscriminator, self).__init__()
# we don't want to apply ReLU activation to raw image before convolution transformation.
self.model = nn.Sequential(
SpectralNorm(nn.Conv2d(in_channels, out_channels, 3, 1,
padding=1)), nn.ReLU(),
SpectralNorm(nn.Conv2d(out_channels, out_channels, 3, 1,
padding=1)), nn.AvgPool2d(2))
self.bypass = nn.Sequential(
nn.AvgPool2d(2),
SpectralNorm(nn.Conv2d(in_channels, out_channels, 1, 1,
padding=0)),
)
def Linear(in_features, out_features, sn_norm):
if sn_norm:
linear = SpectralNorm(nn.Linear(in_features, out_features))
else:
linear = nn.Linear(in_features, out_features)
return linear
def __init__(self):
super(DiscriminatorRes, self).__init__()
self.model = nn.Sequential(
# 128 x 32 x 32
FirstResBlockDiscriminator(3, 64, stride=2),
# 128 x 16 x 16
ResBlockDiscriminator(64, 128, stride=2),
# 128 x 8 x 8
ResBlockDiscriminator(128, 256, stride=2),
# 128 x 4 x 4
ResBlockDiscriminator(256, 512, stride=2),
# 128 x 4 x 4
ResBlockDiscriminator(512, 512, stride=1),
ResBlockDiscriminator(512, 512, stride=1),
nn.ReLU(),
nn.AvgPool2d(4)
)
self.fc = SpectralNorm(nn.Linear(512, 1))
def Conv2d(in_channels, out_channels, kernel=3, stride=1, padding=1, sn_norm=False):
if sn_norm:
conv = SpectralNorm(nn.Conv2d(in_channels, out_channels, kernel, stride, padding))
else:
conv = nn.Conv2d(in_channels, out_channels, kernel, stride, padding)
return conv
def __init__(self, in_channels, out_channels, stride=1):
super(ResBlockDiscriminator, self).__init__()
self.model = nn.Sequential(
nn.ReLU(),
SpectralNorm(
nn.Conv2d(in_channels, out_channels, 3, stride, padding=1)),
nn.ReLU(),
SpectralNorm(nn.Conv2d(out_channels, out_channels, 3, 1,
padding=1)))
self.bypass = nn.Sequential()
if stride != 1:
self.bypass = nn.Sequential(
SpectralNorm(
nn.Conv2d(in_channels, out_channels, 3, stride,
padding=1)))
def __init__(self, z_dim):
super(Generator, self).__init__()
self.z_dim = z_dim
self.conv1 = SpectralNorm(nn.ConvTranspose2d(z_dim, 512, 4, stride=1))
self.bn1 = ConditionalBatchNorm2d(512, num_classes)
self.conv2 = SpectralNorm(
nn.ConvTranspose2d(512, 256, 4, stride=2, padding=(1, 1)))
self.bn2 = ConditionalBatchNorm2d(256, num_classes)
self.conv3 = SpectralNorm(
nn.ConvTranspose2d(256, 128, 4, stride=2, padding=(2, 2)))
self.bn3 = ConditionalBatchNorm2d(128, num_classes)
self.conv4 = SpectralNorm(
nn.ConvTranspose2d(128, 64, 4, stride=2, padding=(1, 1)))
self.bn4 = ConditionalBatchNorm2d(64, num_classes)
self.conv5 = SpectralNorm(
nn.Conv2d(64, channels, 3, stride=1, padding=(1, 1)))