What are the drawbacks/advantages of defining a function that builds a given layer as a member of the neural network class as opposed to defining it outside the class?
For example, as a class member:
class Discriminator(nn.Module):
"""C64-C128-C256-C512 PatchGAN Discriminator architecture."""
def __init__(self, in_channels=3):
super(Discriminator, self).__init__()
channels = [in_channels, 64, 128, 256, 512, 1]
layers = len(channels)
channels = zip(channels, channels[1:])
self.blocks = nn.ModuleList()
for layer, (input_size, output_size) in enumerate(channels):
if layer == 0:
input_size *= 2
batch_norm = False
leaky = True
elif layer < layers - 2:
batch_norm = True
leaky = True
else:
batch_norm = False
leaky = False
self.blocks.append(self.make_conv(in_size=input_size,
out_size=output_size,
batch_norm=batch_norm,
leaky=leaky))
self.init_weights(mean=0.0, std=0.02)
def make_conv(self, in_size, out_size, batch_norm, leaky):
"""Convolutional blocks of the Discriminator.
Let Ck denote a Convolution-BtachNorm-ReLU block with k channels.
All convolutions are 4 x 4 spatial filters with stride 2 and
downsample by a factor of 2. BatchNorm is not applied to the c64 block.
After the C512 block, a convolution is applied to map to a 1-d output,
followed by a Sigmoid function. All ReLUs are leaky with slope of 0.2.
"""
block = [nn.Conv2d(in_size, out_size,
kernel_size=4, stride=2, padding=2,
padding_mode="reflect",
bias=False if batch_norm else True)]
if batch_norm:
block.append(nn.BatchNorm2d(out_size))
if leaky:
block.append(nn.LeakyReLU(0.2))
else:
block.append(nn.Sigmoid())
return nn.Sequential(*block)
def init_weights(self, mean=0.0, std=0.02):
"""Initialize weights from a Gaussian distribution."""
for module in self.modules():
if isinstance(module, (nn.Conv2d, nn.BatchNorm2d)):
nn.init.normal_(module.weight.data, mean=mean, std=std)
def forward(self, x, y):
"""Output an nxn tensor belonging to patch ij in the input image."""
x = torch.cat((x, y), dim=1)
for block in self.blocks:
x = block(x)
return x
outside the class:
def make_conv(self, in_size, out_size, batch_norm, leaky):
"""Convolutional blocks of the Discriminator.
Let Ck denote a Convolution-BtachNorm-ReLU block with k channels.
All convolutions are 4 x 4 spatial filters with stride 2 and
downsample by a factor of 2. BatchNorm is not applied to the c64 block.
After the C512 block, a convolution is applied to map to a 1-d output,
followed by a Sigmoid function. All ReLUs are leaky with slope of 0.2.
"""
block = [nn.Conv2d(in_size, out_size,
kernel_size=4, stride=2, padding=2,
padding_mode="reflect",
bias=False if batch_norm else True)]
if batch_norm:
block.append(nn.BatchNorm2d(out_size))
if leaky:
block.append(nn.LeakyReLU(0.2))
else:
block.append(nn.Sigmoid())
return nn.Sequential(*block)
class Discriminator(nn.Module):
"""C64-C128-C256-C512 PatchGAN Discriminator architecture."""
def __init__(self, in_channels=3):
super(Discriminator, self).__init__()
channels = [in_channels, 64, 128, 256, 512, 1]
layers = len(channels)
channels = zip(channels, channels[1:])
self.blocks = nn.ModuleList()
for layer, (input_size, output_size) in enumerate(channels):
if layer == 0:
input_size *= 2
batch_norm = False
leaky = True
elif layer < layers - 2:
batch_norm = True
leaky = True
else:
batch_norm = False
leaky = False
self.blocks.append(make_conv(in_size=input_size,
out_size=output_size,
batch_norm=batch_norm,
leaky=leaky))
self.init_weights(mean=0.0, std=0.02)
def init_weights(self, mean=0.0, std=0.02):
"""Initialize weights from a Gaussian distribution."""
for module in self.modules():
if isinstance(module, (nn.Conv2d, nn.BatchNorm2d)):
nn.init.normal_(module.weight.data, mean=mean, std=std)
def forward(self, x, y):
"""Output an nxn tensor belonging to patch ij in the input image."""
x = torch.cat((x, y), dim=1)
for block in self.blocks:
x = block(x)
return x