model.py

import torch
import torch.nn as nn
from torchvision.utils import save_image
from torchvision.utils import make_grid
import matplotlib

matplotlib.use('TkAgg')
import matplotlib.pyplot as plt

def weights_init_general(model, mean, std):
    for m in model._modules:
        if isinstance(m, nn.ConvTranspose2d) or isinstance(m, nn.Conv2d):
            model._modules[m].weight.data.normal_(mean, std)
            model._modules[m].bias.data.zero_()


class Generator(nn.Module):
    def __init__(self, input_size, general_complexity, dropout_prob, weights_mean, weights_std, image_channels):
        super(Generator, self).__init__()

        self.loss = nn.BCELoss()

        self.layer1 = nn.Sequential(
            nn.ConvTranspose2d(input_size, 4 * general_complexity, 4, 1, 0, bias=False),
            nn.BatchNorm2d(4 * general_complexity),
            nn.ReLU(True),
            nn.Dropout2d(dropout_prob)
        )
        self.layer2 = nn.Sequential(
            nn.ConvTranspose2d(4 * general_complexity, 2 * general_complexity, 4, 2, 1, bias=False),
            nn.BatchNorm2d(2 * general_complexity),
            nn.ReLU(True),
            nn.Dropout2d(dropout_prob)
        )
        self.layer3 = nn.Sequential(
            nn.ConvTranspose2d(2 * general_complexity, 1 * general_complexity, 4, 2, 1, bias=False),
            nn.BatchNorm2d(1 * general_complexity),
            nn.ReLU(True),
            nn.Dropout2d(dropout_prob)
        )
        self.layer4 = nn.Sequential(
            nn.ConvTranspose2d(1 * general_complexity, 1 * image_channels, 4, 2, 1, bias=False),
            nn.Tanh()
        )

        self.all_layers = nn.Sequential(
            self.layer1,
            self.layer2,
            self.layer3,
            self.layer4)

        weights_init_general(self, weights_mean, weights_std)

    def forward(self, input):
        output = self.all_layers(input)
        return output


class Discriminator(nn.Module):
    def __init__(self, general_complexity, weights_mean, weights_std, packing, image_channels):
        super(Discriminator, self).__init__()

        self.loss = nn.BCELoss()

        self.layer1 = nn.Sequential(
            nn.Conv2d(image_channels * packing, general_complexity, 4, 2, 1, bias=False),
            nn.LeakyReLU(0.2, inplace=True)
        )
        self.layer2 = nn.Sequential(
            nn.Conv2d(general_complexity, general_complexity * 2, 4, 2, 1, bias=False),
            nn.BatchNorm2d(general_complexity * 2),
            nn.LeakyReLU(0.2, inplace=True)
        )
        self.layer3 = nn.Sequential(
            nn.Conv2d(general_complexity * 2, general_complexity * 4, 4, 2, 1, bias=False),
            nn.BatchNorm2d(general_complexity * 4),
            nn.LeakyReLU(0.2, inplace=True)
        )
        self.layer4 = nn.Sequential(
            nn.Conv2d(general_complexity * 4, 1, 4, 1, 0, bias=False),
            nn.Sigmoid()
        )

        self.all_layers = nn.Sequential(
            self.layer1,
            self.layer2,
            self.layer3,
            self.layer4
        )

        weights_init_general(self, weights_mean, weights_std)

    def forward(self, input):
        output = self.all_layers(input)
        return output


def pack(input, packing):
    # Number of elements that need to be added to the input tensor
    nb_to_add = (packing - (input.shape[0] % packing)) % packing

    # Add elements to the input if not a round number for the packing number
    if nb_to_add > 0:
        input = torch.cat((input, input[-nb_to_add:].view(nb_to_add, 3, input.shape[2], input.shape[3])))

    # Reshape the tensor so it is packed
    packed_output = input.view(-1, input.shape[1] * packing, input.shape[2], input.shape[3])
    return packed_output


# Initialise weights of the model with certain mean and standard deviation
def weights_init_general(model, mean, std):
    for m in model._modules:
        if isinstance(m, nn.ConvTranspose2d) or isinstance(m, nn.Conv2d):
            model._modules[m].weight.data.normal_(mean, std)
            model._modules[m].bias.data.zero_()


def write_loss_plot(loss, loss_label, save_path, clear_plot=True):
    # Plot losses
    plt.plot(loss, label=loss_label)

    plt.legend(loc="best")
    plt.xlabel("Epoch")
    plt.ylabel("Loss")

    plt.savefig(save_path + "losses.png")

    if clear_plot:
        plt.clf()


def save_model(model, save_path, name):
    print("Saving " + name + " to : " + save_path)
    torch.save(model.state_dict(), save_path + name + ".pt")


def save_parameters(save_path, file_name="hyperparameters.py"):
    from shutil import copyfile
    copyfile("hyperparameters.py", save_path + "" + file_name)


def rescale_for_rgb_plot(images):
    min_val = images.data.min()
    max_val = images.data.max()
    return (images.data - min_val) / (max_val - min_val)


def save_images(data, save_path, image_size, image_channels, num_row, epoch):
    image_list = []
    for i in range(len(data)):
        image_data = data[i].view(image_channels, image_size, image_size)
        image_data = rescale_for_rgb_plot(image_data)
        image_list.append(image_data)
    save_image(make_grid(image_list, nrow=num_row), save_path + "epoch_" + str(epoch) + ".png")