from torch.nn import Module, Linear, ReLU, Sequential
from torch.nn.utils import spectral_norm
from torchinfo import summary
from collections import OrderedDict as OrderedDict
import torch
import torch.nn as nn
[docs]
class BoundedActivation(Module):
"""
Bounded Activation Layer that applies a bounded activation function using tanh.
Methods:
bounded_activation(x): Apply the bounded activation to the input tensor.
forward(input): Forward pass to apply bounded activation.
"""
def __init__(self, M=100.0):
super().__init__()
self.M = M
[docs]
def forward(self, input):
"""
Forward pass through the bounded activation layer.
Args:
input (torch.Tensor): Input tensor.
Returns:
torch.Tensor: Output after applying bounded activation.
"""
return self.M * torch.tanh(input / self.M)
[docs]
class Discriminator(Module):
"""
Discriminator Class that initializes and processes the discriminator network.
Args:
input_dim (int): Dimension of the input features.
batch_size (int): Batch size for processing.
spec_norm (bool): Whether to apply spectral normalization to the layers.
bounded (bool): Whether to apply bounded activation on the final output.
layers_list (list): List of integers specifying the number of units for each hidden layer.
device (str): Device to run the model on ('cpu' or 'cuda').
Methods:
forward(inputs): Forward pass through the network.
Returns:
torch.Tensor: Output after processing through the discriminator network.
"""
def __init__(self, input_dim, batch_size, spec_norm, bounded, layers_list, device='cuda' if torch.cuda.is_available() else 'cpu'):
super(Discriminator, self).__init__()
self.input_dim = input_dim
self.spec_norm = spec_norm
self.layers_list = layers_list
self.bounded = bounded
self.batch_size = batch_size
model_dict = OrderedDict()
if self.spec_norm:
for i, h_dim in enumerate(self.layers_list):
model_dict[f'Dense{i}'] = spectral_norm(Linear(input_dim, h_dim))
model_dict[f'ReLU{i}'] = ReLU()
input_dim = h_dim
model_dict[f'Dense{i+1}'] = spectral_norm(Linear(input_dim, 1))
else:
for i, h_dim in enumerate(self.layers_list):
model_dict[f'Dense{i}'] = Linear(input_dim, h_dim)
model_dict[f'ReLU{i}'] = ReLU()
input_dim = h_dim
model_dict[f'Dense{i+1}'] = Linear(input_dim, 1)
if bounded:
model_dict[f'Bounded_Activation'] = BoundedActivation()
self.discriminator = Sequential(model_dict).to(device)
# print()
# print('Discriminator Summary:')
# summary(self.discriminator, (self.batch_size, self.input_dim))
[docs]
def forward(self, inputs):
"""
Forward pass through the discriminator.
Args:
inputs (torch.Tensor): Input data to the network.
Returns:
torch.Tensor: Discriminator's prediction after processing the inputs.
"""
predicted = self.discriminator(inputs)
return predicted
[docs]
class Generator(Module):
"""
Generator Class that initializes and processes the generator network.
Args:
X_dim (int): Dimension of the output generated by the generator.
Z_dim (int): Dimension of the input latent space.
batch_size (int): Batch size for processing.
spec_norm (bool): Whether to apply spectral normalization to the layers.
layers_list (list): List of integers specifying the number of units for each hidden layer.
device (str): Device to run the model on ('cpu' or 'cuda').
Methods:
forward(inputs): Forward pass through the network.
Returns:
torch.Tensor: Generated output after processing the latent input.
"""
def __init__(self, X_dim, Z_dim, batch_size, spec_norm, layers_list, device='cuda' if torch.cuda.is_available() else 'cpu'):
super(Generator, self).__init__()
self.X_dim = X_dim
self.Z_dim = Z_dim
self.batch_size = batch_size
self.spec_norm = spec_norm
self.layers_list = layers_list
model_dict = OrderedDict()
input_dim = self.Z_dim
if spec_norm:
for i, h_dim in enumerate(self.layers_list):
model_dict[f'Dense{i}'] = spectral_norm(Linear(input_dim, h_dim))
model_dict[f'ReLU{i}'] = ReLU()
input_dim = h_dim
model_dict[f'Dense{i+1}'] = spectral_norm(Linear(input_dim, self.X_dim))
else:
for i, h_dim in enumerate(self.layers_list):
model_dict[f'Dense{i}'] = Linear(input_dim, h_dim)
model_dict[f'ReLU{i}'] = ReLU()
input_dim = h_dim
model_dict[f'Dense{i+1}'] = Linear(input_dim, X_dim)
self.generator = Sequential(model_dict).to(device)
print()
print('Generator Summary:')
summary(self.generator, (self.batch_size, self.Z_dim))
[docs]
def forward(self, inputs):
"""
Forward pass through the generator.
Args:
inputs (torch.Tensor): Input latent vector.
Returns:
torch.Tensor: Generated output.
"""
generated = self.generator(inputs)
return generated
[docs]
class Discriminator_MNIST(nn.Module):
"""
Discriminator Class for MNIST that processes the input and classifies real vs fake images.
Methods:
forward(x): Forward pass through the network.
Returns:
torch.Tensor: Output after processing the input image through the discriminator network.
"""
def __init__(self):
super(Discriminator_MNIST, self).__init__()
self.linear0 = nn.Linear(784, 794)
self.linear1 = nn.Linear(794, 794)
self.linear2 = nn.Linear(794, 256)
self.linear3 = nn.Linear(256, 128)
self.linear4 = nn.Linear(128, 64)
self.linear5 = nn.Linear(64, 32)
self.linear6 = nn.Linear(32, 16)
self.linear7 = nn.Linear(16, 1)
self.ln0 = nn.LayerNorm(794)
self.ln1 = nn.LayerNorm(794)
self.ln2 = nn.LayerNorm(256)
self.ln3 = nn.LayerNorm(128)
self.ln4 = nn.LayerNorm(64)
self.ln5 = nn.LayerNorm(32)
self.relu = nn.LeakyReLU(0.2)
[docs]
def forward(self, x):
"""
Forward pass through the MNIST discriminator.
Args:
x (torch.Tensor): Input image data reshaped to (batch_size, 784).
Returns:
torch.Tensor: Output after processing through dense and normalization layers.
"""
x = x.reshape(-1, 784)
w = self.relu(self.ln0(self.linear0(x)))
w = self.relu(self.ln1(self.linear1(w)))
w = self.relu(self.ln2(self.linear2(w)))
w = self.relu(self.ln3(self.linear3(w)))
w = self.relu(self.ln4(self.linear4(w)))
w = self.relu(self.ln5(self.linear5(w)))
w = self.relu(self.linear6(w))
out = self.linear7(w)
return out