def __init__(self, img_dim):
super().__init__()
self.disc = nn.Sequential(nn.Linear(img_dim, 128), nn.LeakyReLU(0.1),
nn.Linear(128, 1), nn.Sigmoid())
def __init__(self, backbone):
super().__init__()
self.feature_extractor, num_features = self.adapt_backbone(backbone)
self.fc_age = nn.Sequential(nn.Linear(num_features, 7),
nn.LeakyReLU(0.2))
def block(in_feat, out_feat, normalize=True):
layers = [nn.Linear(in_feat, out_feat)]
if normalize:
layers.append(nn.BatchNorm1d(out_feat, 0.8))
layers.append(nn.LeakyReLU(0.2, inplace=True))
return layers
def __init__(self, nz, n_label=0, arch=None):
super().__init__()
self.nz = nz
#self.tensor_properties = torch.ones(1).to(device=args.device) #hack
if n_label > 0:
self.label_embed = nn.Embedding(n_label, n_label)
self.label_embed.weight.data.normal_()
self.code_norm = PixelNorm()
self.adanorm_blocks = nn.ModuleList()
#self.z_const = torch.ones(512, 4, 4).to(device=args.device)
HLM = 1 if arch == 'small' else 2 # High-resolution Layer multiplier: Use to make the 64x64+ resolution layers larger by this factor (1 = default Balanced Pioneer)
progression_raw = [
ConvBlock(nz,
nz,
4,
3,
3,
1,
spectral_norm=gen_spectral_norm,
const_layer=True,
holder=self),
ConvBlock(nz,
nz,
3,
1,
spectral_norm=gen_spectral_norm,
holder=self),
ConvBlock(nz,
nz,
3,
1,
spectral_norm=gen_spectral_norm,
holder=self),
ConvBlock(nz,
nz,
3,
1,
spectral_norm=gen_spectral_norm,
holder=self),
ConvBlock(nz,
int(nz / 2) * HLM,
3,
1,
spectral_norm=gen_spectral_norm,
holder=self),
ConvBlock(int(nz / 2) * HLM,
int(nz / 4) * HLM,
3,
1,
spectral_norm=gen_spectral_norm,
holder=self),
ConvBlock(int(nz / 4) * HLM,
int(nz / 8) * HLM,
3,
1,
spectral_norm=gen_spectral_norm,
holder=self),
ConvBlock(int(nz / 8) * HLM,
int(nz / 16) * HLM,
3,
1,
spectral_norm=gen_spectral_norm,
holder=self),
ConvBlock(int(nz / 16) * HLM,
int(nz / 32) * HLM,
3,
1,
spectral_norm=gen_spectral_norm,
holder=self)
]
to_rgb_raw = [
nn.Conv2d(nz, 3, 1), #Each has 3 out channels and kernel size 1x1!
nn.Conv2d(nz, 3, 1),
nn.Conv2d(nz, 3, 1),
nn.Conv2d(nz, 3, 1),
nn.Conv2d(int(nz / 2) * HLM, 3, 1),
nn.Conv2d(int(nz / 4) * HLM, 3, 1),
nn.Conv2d(int(nz / 8) * HLM, 3, 1),
nn.Conv2d(int(nz / 16) * HLM, 3, 1),
nn.Conv2d(int(nz / 32) * HLM, 3, 1)
]
noise_raw = [
nn.ModuleList([NoiseLayer(nz), NoiseLayer(nz)]),
nn.ModuleList([NoiseLayer(nz), NoiseLayer(nz)]),
nn.ModuleList([NoiseLayer(nz), NoiseLayer(nz)]),
nn.ModuleList([NoiseLayer(nz), NoiseLayer(nz)]),
nn.ModuleList(
[NoiseLayer(int(nz / 2) * HLM),
NoiseLayer(int(nz / 2) * HLM)]),
nn.ModuleList(
[NoiseLayer(int(nz / 4) * HLM),
NoiseLayer(int(nz / 4) * HLM)]),
nn.ModuleList(
[NoiseLayer(int(nz / 8) * HLM),
NoiseLayer(int(nz / 8) * HLM)]),
nn.ModuleList([
NoiseLayer(int(nz / 16) * HLM),
NoiseLayer(int(nz / 16) * HLM)
]),
nn.ModuleList([
NoiseLayer(int(nz / 32) * HLM),
NoiseLayer(int(nz / 32) * HLM)
])
]
# The args.flip_invariance_layer (when >=0) relates to driving scale-specific invariances in the weakly supervised case.
# We disable noise by default if there are support layers. This is only to replicate the approach in Deep Automodulators paper.
# Feel free to enable it otherwise.
self.has_noise_layers = args.flip_invariance_layer <= -1
if args.flip_invariance_layer <= -1:
self.progression = nn.ModuleList(progression_raw)
self.to_rgb = nn.ModuleList(to_rgb_raw)
if self.has_noise_layers:
self.noise = nn.ModuleList(noise_raw)
Generator.supportBlockPoints = []
else:
self.supportTorgbBlock8 = nn.Conv2d(nz, 3, 1)
self.supportProgression = ConvBlock(
nz, nz, 3, 1, spectral_norm=gen_spectral_norm, holder=self)
if self.has_noise_layers:
self.supportNoise = nn.ModuleList(
[NoiseLayer(nz), NoiseLayer(nz)])
# Add support blocks like this and also to the Generator.supportBlockPoints
#TODO: Support adding multiple layers in the following ModuleList concatenators:
self.progression = nn.ModuleList(
progression_raw[:args.flip_invariance_layer] +
[self.supportProgression] +
progression_raw[args.flip_invariance_layer:])
self.to_rgb = nn.ModuleList(
to_rgb_raw[:args.flip_invariance_layer] +
[self.supportTorgbBlock8] +
to_rgb_raw[args.flip_invariance_layer:])
if self.has_noise_layers:
self.noise = nn.ModuleList(
noise_raw[:args.flip_invariance_layer] +
[self.supportNoise] +
noise_raw[args.flip_invariance_layer:])
Generator.supportBlockPoints = [
args.flip_invariance_layer
] # You can add several layers here as a list, but you need to add the argparser support for that.
mapping_lrmul = 0.01
self.use_layer_noise = (not args.no_LN
and args.flip_invariance_layer <= -1)
self.z_preprocess = nn.Sequential(
nn.Sequential(equal_lr(nn.Linear(nz, nz), gain=mapping_lrmul),
nn.LeakyReLU(0.2)),
nn.Sequential(equal_lr(nn.Linear(nz, nz), gain=mapping_lrmul),
nn.LeakyReLU(0.2)))
init_linear(self.z_preprocess[0][0], lr_gain=mapping_lrmul)
init_linear(self.z_preprocess[1][0], lr_gain=mapping_lrmul)
def __init__(
self,
in_channels: int,
latent_dim: int,
hidden_dims: List = None,
alpha: float = 100.0,
beta: float = 10.0,
lr: float = 0.005,
weight_decay: Optional[float] = 0,
scheduler_gamma: Optional[float] = 0.97,
) -> None:
super(LogCoshVAE, self).__init__(
lr=lr, weight_decay=weight_decay, scheduler_gamma=scheduler_gamma
)
self.latent_dim = latent_dim
self.alpha = alpha
self.beta = beta
modules = []
if hidden_dims is None:
hidden_dims = [32, 64, 128, 256, 512]
# Build Encoder
for h_dim in hidden_dims:
modules.append(
nn.Sequential(
nn.Conv2d(
in_channels,
out_channels=h_dim,
kernel_size=3,
stride=2,
padding=1,
),
nn.BatchNorm2d(h_dim),
nn.LeakyReLU(),
)
)
in_channels = h_dim
self.encoder = nn.Sequential(*modules)
self.fc_mu = nn.Linear(hidden_dims[-1] * 4, latent_dim)
self.fc_var = nn.Linear(hidden_dims[-1] * 4, latent_dim)
# Build Decoder
modules = []
self.decoder_input = nn.Linear(latent_dim, hidden_dims[-1] * 4)
hidden_dims.reverse()
for i in range(len(hidden_dims) - 1):
modules.append(
nn.Sequential(
nn.ConvTranspose2d(
hidden_dims[i],
hidden_dims[i + 1],
kernel_size=3,
stride=2,
padding=1,
output_padding=1,
),
nn.BatchNorm2d(hidden_dims[i + 1]),
nn.LeakyReLU(),
)
)
self.decoder = nn.Sequential(*modules)
self.final_layer = nn.Sequential(
nn.ConvTranspose2d(
hidden_dims[-1],
hidden_dims[-1],
kernel_size=3,
stride=2,
padding=1,
output_padding=1,
),
nn.BatchNorm2d(hidden_dims[-1]),
nn.LeakyReLU(),
nn.Conv2d(hidden_dims[-1], out_channels=3, kernel_size=3, padding=1),
nn.Tanh(),
)
def __init__(self, in_channels=19, n_classes=15, base_n_filter = 8):
super(UGenerator, self).__init__()
self.in_channels = in_channels
self.n_classes = n_classes
self.base_n_filter = base_n_filter
self.lrelu = nn.LeakyReLU()
self.dropout2d = nn.Dropout2d(p=0.3)
self.upsacle = nn.Upsample(scale_factor=2, mode='nearest')
self.softmax = nn.Softmax(dim=1)
# Level 1 context pathway
self.conv2d_c1_1 = nn.Conv2d(self.in_channels, self.base_n_filter, kernel_size=3, stride=1, padding=1, bias=False)
self.conv2d_c1_2 = nn.Conv2d(self.base_n_filter, self.base_n_filter, kernel_size=3, stride=1, padding=1, bias=False)
self.lrelu_conv_c1 = self.lrelu_conv(self.base_n_filter, self.base_n_filter)
self.inorm2d_c1 = nn.InstanceNorm2d(self.base_n_filter)
# Level 2 context pathway
self.conv2d_c2 = nn.Conv2d(self.base_n_filter, self.base_n_filter*2, kernel_size=3, stride=2, padding=1, bias=False)
self.norm_lrelu_conv_c2 = self.norm_lrelu_conv(self.base_n_filter*2, self.base_n_filter*2)
self.inorm2d_c2 = nn.InstanceNorm2d(self.base_n_filter*2)
# Level 3 context pathway
self.conv2d_c3 = nn.Conv2d(self.base_n_filter*2, self.base_n_filter*4, kernel_size=3, stride=2, padding=1, bias=False)
self.norm_lrelu_conv_c3 = self.norm_lrelu_conv(self.base_n_filter*4, self.base_n_filter*4)
self.inorm2d_c3 = nn.InstanceNorm2d(self.base_n_filter*4)
# Level 4 context pathway
self.conv2d_c4 = nn.Conv2d(self.base_n_filter*4, self.base_n_filter*8, kernel_size=3, stride=2, padding=1, bias=False)
self.norm_lrelu_conv_c4 = self.norm_lrelu_conv(self.base_n_filter*8, self.base_n_filter*8)
self.inorm2d_c4 = nn.InstanceNorm2d(self.base_n_filter*8)
# Level 5 context pathway, level 0 localization pathway
self.conv2d_c5 = nn.Conv2d(self.base_n_filter*8, self.base_n_filter*16, kernel_size=3, stride=2, padding=1, bias=False)
self.norm_lrelu_conv_c5 = self.norm_lrelu_conv(self.base_n_filter*16, self.base_n_filter*16)
self.norm_lrelu_upscale_conv_norm_lrelu_l0 = self.norm_lrelu_upscale_conv_norm_lrelu(self.base_n_filter*16, self.base_n_filter*8)
self.conv2d_l0 = nn.Conv2d(self.base_n_filter*8, self.base_n_filter*8, kernel_size = 1, stride=1, padding=0, bias=False)
self.inorm2d_l0 = nn.InstanceNorm2d(self.base_n_filter*8)
# Level 1 localization pathway
self.conv_norm_lrelu_l1 = self.conv_norm_lrelu(self.base_n_filter*16, self.base_n_filter*16)
self.conv2d_l1 = nn.Conv2d(self.base_n_filter*16, self.base_n_filter*8, kernel_size=1, stride=1, padding=0, bias=False)
self.norm_lrelu_upscale_conv_norm_lrelu_l1 = self.norm_lrelu_upscale_conv_norm_lrelu(self.base_n_filter*8, self.base_n_filter*4)
# Level 2 localization pathway
self.conv_norm_lrelu_l2 = self.conv_norm_lrelu(self.base_n_filter*8, self.base_n_filter*8)
self.conv2d_l2 = nn.Conv2d(self.base_n_filter*8, self.base_n_filter*4, kernel_size=1, stride=1, padding=0, bias=False)
self.norm_lrelu_upscale_conv_norm_lrelu_l2 = self.norm_lrelu_upscale_conv_norm_lrelu(self.base_n_filter*4, self.base_n_filter*2)
# Level 3 localization pathway
self.conv_norm_lrelu_l3 = self.conv_norm_lrelu(self.base_n_filter*4, self.base_n_filter*4)
self.conv2d_l3 = nn.Conv2d(self.base_n_filter*4, self.base_n_filter*2, kernel_size=1, stride=1, padding=0, bias=False)
self.norm_lrelu_upscale_conv_norm_lrelu_l3 = self.norm_lrelu_upscale_conv_norm_lrelu(self.base_n_filter*2, self.base_n_filter)
# Level 4 localization pathway
self.conv_norm_lrelu_l4 = self.conv_norm_lrelu(self.base_n_filter*2, self.base_n_filter*2)
self.conv2d_l4 = nn.Conv2d(self.base_n_filter*2, self.n_classes, kernel_size=1, stride=1, padding=0, bias=False)
self.ds2_1x1_conv2d = nn.Conv2d(self.base_n_filter*8, self.n_classes, kernel_size=1, stride=1, padding=0, bias=False)
self.ds3_1x1_conv2d = nn.Conv2d(self.base_n_filter*4, self.n_classes, kernel_size=1, stride=1, padding=0, bias=False)
def lrelu_conv(self, feat_in, feat_out):
return nn.Sequential(
nn.LeakyReLU(),
nn.Conv2d(feat_in, feat_out, kernel_size=3, stride=1, padding=1, bias=False))
def __init__(self, opt):
super(Feedback, self).__init__()
self.fc1 = nn.Linear(opt.ngh, opt.ngh)
self.fc2 = nn.Linear(opt.ngh, opt.ngh)
self.lrelu = nn.LeakyReLU(0.2, True)
self.apply(weights_init)
def discriminator_block(in_filters, out_filters, bn=True):
"""Returns layers of each discriminator block"""
block = [nn.Conv2d(in_filters, out_filters, 3, 2, 1), nn.LeakyReLU(0.2, inplace=True), nn.Dropout2d(0.25)]
if bn:
block.append(nn.BatchNorm2d(out_filters, 0.8))
return block
def __init__(self,
in_channels,
out_channels,
nf0,
num_down,
max_channels,
use_dropout,
upsampling_mode='transpose',
dropout_prob=0.1,
norm=nn.BatchNorm2d,
outermost_linear=False,
added_channels=2):
'''
:param in_channels: Number of input channels
:param out_channels: Number of output channels
:param nf0: Number of features at highest level of U-Net
:param num_down: Number of downsampling stages.
:param max_channels: Maximum number of channels (channels multiply by 2 with every downsampling stage)
:param use_dropout: Whether to use dropout or no.
:param dropout_prob: Dropout probability if use_dropout=True.
:param upsampling_mode: Which type of upsampling should be used. See "UpBlock" for documentation.
:param norm: Which norm to use. If None, no norm is used. Default is Batchnorm with affinity.
:param outermost_linear: Whether the output layer should be a linear layer or a nonlinear one.
'''
super().__init__()
assert (num_down > 0), "Need at least one downsampling layer in UNet."
# Define the in block
self.in_layer = [
Conv2dSame(in_channels,
nf0,
kernel_size=3,
bias=True if norm is None else False)
]
if norm is not None:
self.in_layer += [norm(nf0, affine=True)]
self.in_layer += [nn.LeakyReLU(0.2, True)]
if use_dropout:
self.in_layer += [nn.Dropout2d(dropout_prob)]
self.in_layer = nn.Sequential(*self.in_layer)
# Define the center UNet block
self.unet_block = UnetSkipConnectionBlock(
min(2**(num_down - 1) * nf0, max_channels),
min(2**(num_down - 1) * nf0, max_channels),
use_dropout=use_dropout,
dropout_prob=dropout_prob,
norm=None, # Innermost has no norm (spatial dimension 1)
upsampling_mode=upsampling_mode,
level=num_down,
added_channels=added_channels)
for i in list(range(0, num_down - 1))[::-1]:
self.unet_block = UnetSkipConnectionBlock(
min(2**i * nf0, max_channels),
min(2**(i + 1) * nf0, max_channels),
use_dropout=use_dropout,
dropout_prob=dropout_prob,
submodule=self.unet_block,
norm=norm,
upsampling_mode=upsampling_mode,
level=i + 1,
added_channels=added_channels)
# Define the out layer. Each unet block concatenates its inputs with its outputs - so the output layer
# automatically receives the output of the in_layer and the output of the last unet layer.
self.out_layer = [
Conv2dSame(2 * nf0 + added_channels,
out_channels,
kernel_size=3,
bias=outermost_linear or (norm is None))
]
if not outermost_linear:
if norm is not None:
self.out_layer += [norm(out_channels, affine=True)]
self.out_layer += [nn.ReLU(True)]
if use_dropout:
self.out_layer += [nn.Dropout2d(dropout_prob)]
self.out_layer = nn.Sequential(*self.out_layer)
self.out_layer_weight = self.out_layer[0].weight
def __init__(self, opt):
super(Discriminator_D1, self).__init__()
self.fc1 = nn.Linear(opt.resSize + opt.attSize, opt.ndh)
self.fc2 = nn.Linear(opt.ndh, 1)
self.lrelu = nn.LeakyReLU(0.2, True)
self.apply(weights_init)
def __init__(self,
in_channels,
out_channels,
prep_conv=True,
middle_channels=None,
use_dropout=False,
dropout_prob=0.1,
norm=nn.BatchNorm2d):
'''
:param in_channels: Number of input channels
:param out_channels: Number of output channels
:param prep_conv: Whether to have another convolutional layer before the downsampling layer.
:param middle_channels: If prep_conv is true, this sets the number of channels between the prep and downsampling
convs.
:param use_dropout: bool. Whether to use dropout or not.
:param dropout_prob: Float. The dropout probability (if use_dropout is True)
:param norm: Which norm to use. If None, no norm is used. Default is Batchnorm with affinity.
'''
super().__init__()
if middle_channels is None:
middle_channels = in_channels
net = list()
if prep_conv:
net += [
nn.ReflectionPad2d(1),
nn.Conv2d(in_channels,
middle_channels,
kernel_size=3,
padding=0,
stride=1,
bias=True if norm is None else False)
]
if norm is not None:
net += [norm(middle_channels, affine=True)]
net += [nn.LeakyReLU(0.2, True)]
if use_dropout:
net += [nn.Dropout2d(dropout_prob, False)]
net += [
nn.ReflectionPad2d(1),
nn.Conv2d(middle_channels,
out_channels,
kernel_size=4,
padding=0,
stride=2,
bias=True if norm is None else False)
]
if norm is not None:
net += [norm(out_channels, affine=True)]
net += [nn.LeakyReLU(0.2, True)]
if use_dropout:
net += [nn.Dropout2d(dropout_prob, False)]
self.net = nn.Sequential(*net)
def create_conv(in_channels,
out_channels,
kernel_size,
order,
num_groups,
padding=1):
"""
Create a list of modules with together constitute a single conv layer with non-linearity
and optional batchnorm/groupnorm.
*https://github.com/wolny/pytorch-3dunet/blob/master/pytorch3dunet/unet3d/model.py*
Args:
in_channels (int): number of input channels
out_channels (int): number of output channels
order (string): order of things, e.g.
'cr' -> conv + ReLU
'crg' -> conv + ReLU + groupnorm
'cl' -> conv + LeakyReLU
'ce' -> conv + ELU
num_groups (int): number of groups for the GroupNorm
padding (int): add zero-padding to the input
Return:
list of tuple (name, module)
"""
assert 'c' in order, "Conv layer MUST be present"
assert order[
0] not in 'rle', 'Non-linearity cannot be the first operation in the layer'
modules = []
for i, char in enumerate(order):
if char == 'r':
modules.append(('ReLU', nn.ReLU(inplace=True)))
elif char == 'l':
modules.append(
('LeakyReLU', nn.LeakyReLU(negative_slope=0.1, inplace=True)))
elif char == 'e':
modules.append(('ELU', nn.ELU(inplace=True)))
elif char == 'c':
# add learnable bias only in the absence of gatchnorm/groupnorm
bias = not ('g' in order or 'b' in order)
modules.append(('conv',
conv3d(in_channels,
out_channels,
kernel_size,
bias,
padding=padding)))
elif char == 'g':
is_before_conv = i < order.index('c')
assert not is_before_conv, 'GroupNorm MUST go after the Conv3d'
# number of groups must be less or equal the number of channels
if out_channels < num_groups:
num_groups = out_channels
modules.append(('groupnorm',
nn.GroupNorm(num_groups=num_groups,
num_channels=out_channels)))
elif char == 'b':
is_before_conv = i < order.index('c')
if is_before_conv:
modules.append(('batchnorm', nn.BatchNorm3d(in_channels)))
else:
modules.append(('batchnorm', nn.BatchNorm3d(out_channels)))
else:
raise ValueError('error')
return modules
def __init__(self, z_dim, img_dim):
super(Generator, self).__init__()
self.gen = nn.Sequential(nn.Linear(z_dim, 256), nn.LeakyReLU(0.1),
nn.Linear(256, img_dim), nn.Tanh())
def create_modules(blocks):
net_info = blocks[
0] #Captures the information about the input and pre-processing
module_list = nn.ModuleList()
prev_filters = 3
output_filters = []
for index, x in enumerate(blocks[1:]):
module = nn.Sequential()
#check the type of block
#create a new module for the block
#append to module_list
#If it's a convolutional layer
if (x["type"] == "convolutional"):
#Get the info about the layer
activation = x["activation"]
try:
batch_normalize = int(x["batch_normalize"])
bias = False
except:
batch_normalize = 0
bias = True
filters = int(x["filters"])
padding = int(x["pad"])
kernel_size = int(x["size"])
stride = int(x["stride"])
if padding:
pad = (kernel_size - 1) // 2
else:
pad = 0
#Add the convolutional layer
conv = nn.Conv2d(prev_filters,
filters,
kernel_size,
stride,
pad,
bias=bias)
module.add_module("conv_{0}".format(index), conv)
#Add the Batch Norm Layer
if batch_normalize:
bn = nn.BatchNorm2d(filters)
module.add_module("batch_norm_{0}".format(index), bn)
#Check the activation.
#It is either Linear or a Leaky ReLU for YOLO
if activation == "leaky":
activn = nn.LeakyReLU(0.1, inplace=True)
module.add_module("leaky_{0}".format(index), activn)
#If it's an upsampling layer
#We use Bilinear2dUpsampling
elif (x["type"] == "upsample"):
stride = int(x["stride"])
upsample = nn.Upsample(scale_factor=2, mode="nearest")
module.add_module("upsample_{}".format(index), upsample)
#If it is a route layer
elif (x["type"] == "route"):
x["layers"] = x["layers"].split(',')
#Start of a route
start = int(x["layers"][0])
#end, if there exists one.
try:
end = int(x["layers"][1])
except:
end = 0
#Positive anotation
if start > 0:
start = start - index
if end > 0:
end = end - index
route = EmptyLayer()
module.add_module("route_{0}".format(index), route)
if end < 0:
filters = output_filters[index +
start] + output_filters[index + end]
else:
filters = output_filters[index + start]
#shortcut corresponds to skip connection
elif x["type"] == "shortcut":
shortcut = EmptyLayer()
module.add_module("shortcut_{}".format(index), shortcut)
#Yolo is the detection layer
elif x["type"] == "yolo":
mask = x["mask"].split(",")
mask = [int(x) for x in mask]
anchors = x["anchors"].split(",")
anchors = [int(a) for a in anchors]
anchors = [(anchors[i], anchors[i + 1])
for i in range(0, len(anchors), 2)]
anchors = [anchors[i] for i in mask]
detection = DetectionLayer(anchors)
module.add_module("Detection_{}".format(index), detection)
module_list.append(module)
prev_filters = filters
output_filters.append(filters)
return (net_info, module_list)
def __init__(self,
input_nc,
ndf=64,
n_layers=3,
norm_layer=nn.BatchNorm2d):
"""Construct a PatchGAN discriminator
Parameters:
input_nc (int) -- the number of channels in input images
ndf (int) -- the number of filters in the last conv layer
n_layers (int) -- the number of conv layers in the discriminator
norm_layer -- normalization layer
"""
super(NLayerDiscriminator, self).__init__()
if type(
norm_layer
) == functools.partial: # no need to use bias as BatchNorm2d has affine parameters
use_bias = norm_layer.func != nn.BatchNorm2d
else:
use_bias = norm_layer != nn.BatchNorm2d
kw = 4
padw = 1
sequence = [
nn.Conv2d(input_nc, ndf, kernel_size=kw, stride=2, padding=padw),
nn.LeakyReLU(0.2, True)
]
nf_mult = 1
nf_mult_prev = 1
for n in range(1,
n_layers): # gradually increase the number of filters
nf_mult_prev = nf_mult
nf_mult = min(2**n, 8)
sequence += [
nn.Conv2d(ndf * nf_mult_prev,
ndf * nf_mult,
kernel_size=kw,
stride=2,
padding=padw,
bias=use_bias),
norm_layer(ndf * nf_mult),
nn.LeakyReLU(0.2, True)
]
nf_mult_prev = nf_mult
nf_mult = min(2**n_layers, 8)
sequence += [
nn.Conv2d(ndf * nf_mult_prev,
ndf * nf_mult,
kernel_size=kw,
stride=1,
padding=padw,
bias=use_bias),
norm_layer(ndf * nf_mult),
nn.LeakyReLU(0.2, True)
]
sequence += [
nn.Conv2d(ndf * nf_mult, 1, kernel_size=kw, stride=1, padding=padw)
] # output 1 channel prediction map
self.model = nn.Sequential(*sequence)
def __init__(self, Data = 'MNIST', batch_size = 100, z_dim=128): #Data = 'MNIST' or 'CIFAR'
super(CNN_VAE, self).__init__()
if Data == 'MNIST':
self.dim = {'batch_size': batch_size, 'hight':28, 'pixels':784, 'channels':1}
elif Data == 'CIFAR':
self.dim = {'batch_size': batch_size, 'hight':32, 'pixels':1024, 'channels':3}
else: raise ValueError("Data set not support")
self.z_dim = z_dim
if Data == 'MNIST':
self.encoder = nn.Sequential(
# input is (nc) x 28 x 28
nn.Conv2d(1, ndf, 4, 2, 1, bias=False),
nn.LeakyReLU(0.2, inplace=True),
# state size. (ndf) x 14 x 14
nn.Conv2d(ndf, ndf * 2, 4, 2, 1, bias=False),
nn.BatchNorm2d(ndf * 2),
nn.LeakyReLU(0.2, inplace=True),
# state size. (ndf*2) x 7 x 7
nn.Conv2d(ndf * 2, ndf * 4, 3, 2, 1, bias=False),
nn.BatchNorm2d(ndf * 4),
nn.LeakyReLU(0.2, inplace=True),
# state size. (ndf*4) x 4 x 4
flatten(),
nn.Linear(4096,2048),
nn.LeakyReLU(0.2, inplace=True),
nn.Linear(2048,1024),
nn.LeakyReLU(0.2, inplace=True),
nn.Linear(1024,400),
)
elif Data == 'CIFAR':
self.encoder = nn.Sequential(
# input is (nc) x 32 x 32
nn.Conv2d(3, ndf * 2, 4, 2, 1, bias=False),
nn.BatchNorm2d(ndf * 2),
nn.LeakyReLU(0.2, inplace=True),
# state size. (ndf*2) x 16 x 16
nn.Conv2d(ndf * 2, ndf * 4, 4, 2, 1, bias=False),
nn.BatchNorm2d(ndf * 4),
nn.LeakyReLU(0.2, inplace=True),
# state size. (ndf*4) x 8 x 8
nn.Conv2d(ndf * 4, ndf, 4, 2, 1, bias=False),
nn.BatchNorm2d(ndf),
nn.LeakyReLU(0.2, inplace=True),
# state size. (ndf*8) x 4 x 4
flatten(),
nn.Linear(1024,512),
nn.LeakyReLU(0.2, inplace=True),
nn.Linear(512,400),
)
self.fc_mu = nn.Linear(400, self.z_dim)
self.fc_logvar = nn.Linear(400, self.z_dim)
if Data == 'MNIST':
self.decoder = nn.Sequential(
unFlatten(),
# input is Z, going into a convolution
nn.ConvTranspose2d( z_dim, ngf * 8, 4, 1, 0, bias=False),
nn.BatchNorm2d(ngf * 8),
nn.ReLU(True),
# state size. (ngf*8) x 4 x 4
nn.ConvTranspose2d(ngf * 8, ngf * 4, 4, 2, 1, bias=False),
nn.BatchNorm2d(ngf * 4),
nn.ReLU(True),
# state size. (ngf*4) x 8 x 8
nn.ConvTranspose2d( ngf * 4, ngf * 2, 4, 2, 1, bias=False),
nn.BatchNorm2d(ngf * 2),
nn.ReLU(True),
# state size. (ngf*2) x 16 x 16
nn.ConvTranspose2d( ngf * 2, 3, 4, 2, 1, bias=False),
nn.Sigmoid()
)
elif Data == 'CIFAR':
self.decoder = nn.Sequential(
unFlatten(),
# input is Z, going into a convolution
nn.ConvTranspose2d( z_dim, ngf * 8, 4, 1, 0, bias=False),
nn.BatchNorm2d(ngf * 8),
nn.ReLU(True),
# state size. (ngf*8) x 4 x 4
nn.ConvTranspose2d(ngf * 8, ngf * 4, 4, 2, 1, bias=False),
nn.BatchNorm2d(ngf * 4),
nn.ReLU(True),
# state size. (ngf*4) x 8 x 8
nn.ConvTranspose2d( ngf * 4, ngf * 2, 4, 2, 1, bias=False),
nn.BatchNorm2d(ngf * 2),
nn.ReLU(True),
# state size. (ngf*2) x 16 x 16
nn.ConvTranspose2d( ngf * 2, 3, 4, 2, 1, bias=False),
nn.Sigmoid()
)
def make_hidden(structure, normalize):
input_channels=32
layers=[]
output_size = 3000
padding=0
fil_size=5
stride=2
output_size = int((output_size+2*padding-fil_size)/stride)+1
for i in structure:
if i == 'M':
output_size = int((output_size+2*padding-fil_size)/stride)+1
layers += [nn.MaxPool1d(kernel_size=5, stride=2)]
else:
layers += [nn.Conv1d(input_channels, i, kernel_size=3, stride=1, padding=1), nn.LeakyReLU(0.1, inplace=True)]
if normalize:
nn.BatchNorm1d(i, eps=1e-3)
input_channels=i
input_channels=input_channels*output_size
return nn.Sequential(*layers), input_channels, structure
def conv_norm_lrelu(self, feat_in, feat_out):
return nn.Sequential(
nn.Conv2d(feat_in, feat_out, kernel_size=3, stride=1, padding=1, bias=False),
nn.InstanceNorm2d(feat_out),
nn.LeakyReLU())
def generator_block(in_channels, out_channels, bn=True):
layers = [nn.Linear(in_channels, out_channels)]
if bn:
layers.append(nn.BatchNorm1d(out_channels, 0.8))
layers.append(nn.LeakyReLU(0.2, inplace=True))
return layers
def __init__(self,
in_channel,
out_channel,
kernel_size,
padding,
kernel_size2=None,
padding2=None,
pixel_norm=False,
spectral_norm=False,
ada_norm=True,
const_layer=False,
holder=None,
last_act=nn.LeakyReLU(0.2)):
super().__init__()
if last_act is None:
last_act = nn.LeakyReLU(1.0) #NopLayer()
#BLUR HACK
blur = (args.upsampling != 'bilinear')
pad1 = padding
pad2 = padding
if padding2 is not None:
pad2 = padding2
kernel1 = kernel_size
kernel2 = kernel_size
if kernel_size2 is not None:
kernel2 = kernel_size2
self.const_layer = const_layer
if spectral_norm:
if not args.blurSN:
self.conv = nn.Sequential(
SpectralNormConv2d(in_channel,
out_channel,
kernel1,
padding=pad1), nn.LeakyReLU(0.2),
SpectralNormConv2d(out_channel,
out_channel,
kernel2,
padding=pad2), last_act)
else:
self.conv = nn.Sequential(
SpectralNormConv2d(in_channel,
out_channel,
kernel1,
padding=pad1), BlurLayer(),
nn.LeakyReLU(0.2),
SpectralNormConv2d(out_channel,
out_channel,
kernel2,
padding=pad2), last_act)
else:
if ada_norm: # In PGGAN, activation came after PixelNorm. In StyleGAN, PixelNorm/AdaNorm comes after activation.
ada_conv2D = EqualConv2d # nn.Conv2d
#print("AdaNorm layer count: {}".format(len(holder.adanorm_blocks)))
maybeBlur = BlurLayer() if blur else nn.Sequential()
if not const_layer:
if not blur:
firstBlock = nn.Sequential(
ada_conv2D(in_channel,
out_channel,
kernel1,
padding=pad1), nn.LeakyReLU(0.2),
AdaNorm(args.nz + args.n_label,
out_channel,
holder=holder))
else:
firstBlock = nn.Sequential(
ada_conv2D(in_channel,
out_channel,
kernel1,
padding=pad1), BlurLayer(),
nn.LeakyReLU(0.2),
AdaNorm(args.nz + args.n_label,
out_channel,
holder=holder))
else:
firstBlock = nn.Sequential() #Dummy
self.conv = nn.Sequential(
firstBlock,
ada_conv2D(out_channel, out_channel, kernel2,
padding=pad2), nn.LeakyReLU(0.2),
AdaNorm(args.nz + args.n_label, out_channel,
holder=holder))
elif pixel_norm:
self.conv = nn.Sequential(
EqualConv2d(in_channel, out_channel, kernel1,
padding=pad1), PixelNorm(), nn.LeakyReLU(0.2),
EqualConv2d(out_channel,
out_channel,
kernel2,
padding=pad2), PixelNorm(), nn.LeakyReLU(0.2))
else:
self.conv = nn.Sequential(
nn.Conv2d(in_channel, out_channel, kernel1, padding=pad1),
nn.LeakyReLU(0.2),
nn.Conv2d(out_channel, out_channel, kernel2, padding=pad2),
nn.LeakyReLU(0.2))
def conv_feat_block(nIn, nOut):
return nn.Sequential(
nn.Conv2d(nIn, nOut, kernel_size=3, stride=2, padding=1),
nn.LeakyReLU(0.2),
nn.Conv2d(nOut, nOut, kernel_size=3, stride=1, padding=1),
nn.LeakyReLU(0.2))
from typing import Union
import torch
from torch import nn
Activation = Union[str, nn.Module]
_str_to_activation = {
'relu': nn.ReLU(),
'tanh': nn.Tanh(),
'leaky_relu': nn.LeakyReLU(),
'sigmoid': nn.Sigmoid(),
'selu': nn.SELU(),
'softplus': nn.Softplus(),
'identity': nn.Identity(),
}
def build_mlp(
input_size: int,
output_size: int,
n_layers: int,
size: int,
activation: Activation = 'tanh',
output_activation: Activation = 'identity',
):
"""
Builds a feedforward neural network
arguments:
input_placeholder: placeholder variable for the state (batch_size, input_size)
scope: variable scope of the network
def __init__(self, C_in, C_out):
super(Upsample_PS, self).__init__()
self.layer = nn.Sequential(nn.PixelShuffle(2), nn.BatchNorm2d(C_out),
nn.LeakyReLU())
def __init__(self, encoder, decoder, output_activation=nn.LeakyReLU()):
super().__init__()
self.encoder = encoder
self.decoder = decoder
self.out_act = output_activation
def __init__(self, input_features, output_features):
super(_ConvLayer, self).__init__()
self.add_module('conv2', Conv2d(input_features, output_features,
kernel_size=5, stride=2))
self.add_module('leakyrelu', nn.LeakyReLU(0.1, inplace=True))
def __init__(self, classes, device):
super().__init__()
self.classes = classes
self.device = device
self.model_cls = nn.Sequential(
# 64
nn.Conv2d(in_channels=3,
out_channels=128,
kernel_size=4,
stride=2,
padding=1),
nn.LeakyReLU(negative_slope=0.2),
#nn.Dropout2d(p=0.2),
# 32
nn.Conv2d(in_channels=128,
out_channels=256,
kernel_size=4,
stride=2,
padding=1),
nn.BatchNorm2d(256),
nn.LeakyReLU(negative_slope=0.2),
#nn.Dropout2d(p=0.2),
# 16
nn.Conv2d(in_channels=256,
out_channels=512,
kernel_size=4,
stride=2,
padding=1),
nn.BatchNorm2d(512),
nn.LeakyReLU(negative_slope=0.2),
# 8
nn.Conv2d(in_channels=512,
out_channels=512,
kernel_size=3,
stride=1,
padding=1),
nn.BatchNorm2d(512),
nn.LeakyReLU(negative_slope=0.2),
#nn.Dropout2d(p=0.2),
# 8
nn.Conv2d(in_channels=512,
out_channels=1024,
kernel_size=4,
stride=2,
padding=1),
nn.BatchNorm2d(1024),
nn.LeakyReLU(negative_slope=0.2)
# 4
# nn.Linear(1024, classes)
# nn.LeakyReLU(negative_slope=0.2)
)
self.model_decoder = nn.Sequential(
nn.ConvTranspose2d(in_channels=1024,
out_channels=512,
kernel_size=3,
stride=1,
padding=1),
nn.BatchNorm2d(512),
nn.LeakyReLU(negative_slope=0.2),
#nn.UpsamplingNearest2d(scale_factor=2),
# 8
nn.ConvTranspose2d(in_channels=512,
out_channels=256,
kernel_size=4,
stride=2,
padding=1),
nn.BatchNorm2d(256),
nn.LeakyReLU(negative_slope=0.2),
# nn.UpsamplingNearest2d(scale_factor=2),
# 16
nn.ConvTranspose2d(in_channels=256,
out_channels=128,
kernel_size=4,
stride=2,
padding=1),
nn.BatchNorm2d(128),
nn.LeakyReLU(negative_slope=0.2),
# nn.UpsamplingNearest2d(scale_factor=2),
# 32
nn.ConvTranspose2d(in_channels=128,
out_channels=64,
kernel_size=4,
stride=2,
padding=1),
nn.BatchNorm2d(64),
nn.LeakyReLU(negative_slope=0.2),
# nn.UpsamplingNearest2d(scale_factor=2),
# 64
nn.ConvTranspose2d(in_channels=64,
out_channels=32,
kernel_size=4,
stride=2,
padding=1),
nn.BatchNorm2d(32),
nn.LeakyReLU(negative_slope=0.2),
nn.ConvTranspose2d(in_channels=32,
out_channels=3,
kernel_size=3,
stride=1,
padding=1),
# Change to tanh actuvation function
# nn.LeakyReLU(negative_slope=0.2)
nn.Sigmoid())
def __init__(self, input_features, output_features):
super(_UpScale, self).__init__()
self.add_module('conv2_', Conv2d(input_features, output_features * 4,
kernel_size=3))
self.add_module('leakyrelu', nn.LeakyReLU(0.1, inplace=True))
self.add_module('pixelshuffler', _PixelShuffler())
def __init__(self, num_in_ch, num_feat, input_size=128):
super(VGGStyleDiscriminator, self).__init__()
self.input_size = input_size
assert self.input_size == 128 or self.input_size == 256, (
f'input size must be 128 or 256, but received {input_size}')
self.conv0_0 = nn.Conv2d(num_in_ch, num_feat, 3, 1, 1, bias=True)
self.conv0_1 = nn.Conv2d(num_feat, num_feat, 4, 2, 1, bias=False)
self.bn0_1 = nn.BatchNorm2d(num_feat, affine=True)
self.conv1_0 = nn.Conv2d(num_feat, num_feat * 2, 3, 1, 1, bias=False)
self.bn1_0 = nn.BatchNorm2d(num_feat * 2, affine=True)
self.conv1_1 = nn.Conv2d(num_feat * 2,
num_feat * 2,
4,
2,
1,
bias=False)
self.bn1_1 = nn.BatchNorm2d(num_feat * 2, affine=True)
self.conv2_0 = nn.Conv2d(num_feat * 2,
num_feat * 4,
3,
1,
1,
bias=False)
self.bn2_0 = nn.BatchNorm2d(num_feat * 4, affine=True)
self.conv2_1 = nn.Conv2d(num_feat * 4,
num_feat * 4,
4,
2,
1,
bias=False)
self.bn2_1 = nn.BatchNorm2d(num_feat * 4, affine=True)
self.conv3_0 = nn.Conv2d(num_feat * 4,
num_feat * 8,
3,
1,
1,
bias=False)
self.bn3_0 = nn.BatchNorm2d(num_feat * 8, affine=True)
self.conv3_1 = nn.Conv2d(num_feat * 8,
num_feat * 8,
4,
2,
1,
bias=False)
self.bn3_1 = nn.BatchNorm2d(num_feat * 8, affine=True)
self.conv4_0 = nn.Conv2d(num_feat * 8,
num_feat * 8,
3,
1,
1,
bias=False)
self.bn4_0 = nn.BatchNorm2d(num_feat * 8, affine=True)
self.conv4_1 = nn.Conv2d(num_feat * 8,
num_feat * 8,
4,
2,
1,
bias=False)
self.bn4_1 = nn.BatchNorm2d(num_feat * 8, affine=True)
if self.input_size == 256:
self.conv5_0 = nn.Conv2d(num_feat * 8,
num_feat * 8,
3,
1,
1,
bias=False)
self.bn5_0 = nn.BatchNorm2d(num_feat * 8, affine=True)
self.conv5_1 = nn.Conv2d(num_feat * 8,
num_feat * 8,
4,
2,
1,
bias=False)
self.bn5_1 = nn.BatchNorm2d(num_feat * 8, affine=True)
self.linear1 = nn.Linear(num_feat * 8 * 4 * 4, 100)
self.linear2 = nn.Linear(100, 1)
# activation function
self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
def forward(self,x,activation=nn.LeakyReLU(),final_activation=nn.Softmax2d()):
x,tensors,indices,sizes = self.encode(x)
for layer in self.center:
x = activation(layer(x))
x = self.decode((x,tensors,indices,sizes))
return final_activation(self.final(x))