convolutional_nn.py

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.data as data
import torch.optim as optim

from torchvision.datasets import CIFAR10
from torchvision import transforms

import matplotlib.pyplot as plt
import seaborn as sns

# Here we implement an autoencoder using convolutional neural networks

size=32 
base_channel_size=3
hid_channel_size=64
latent_size=16
num_epoch=4
batch_size=100
clip_grad_norm_val=10
lr=0.001


class layer_encoder(nn.Module):
    def __init__(self, size, base_channel_size, hid_channel_size, latent_size):
        super().__init__()
        self.size=size  # (batch_size, base_channel_size, size, size)
        self.base_channel_size=base_channel_size
        self.hid_channel_size=hid_channel_size
        self.latent_size=latent_size
        self.layers=nn.Sequential(
            nn.Conv2d(self.base_channel_size,self.hid_channel_size,kernel_size=3,padding=1,stride=2), # (batch_size, base_channel_size, size, size) -> (batch_size, hid_channel_size, size//2, size//2)
            nn.ReLU(),
        )
        self.conversion_layer=nn.Sequential(
            nn.Flatten(), # (batch_size, hid_channel_size, size//2, size//2) -> (batch_size, hid_channel_size*size//2*size//2)
            nn.Linear(self.hid_channel_size*self.size*self.size//4,self.latent_size) # (batch_size, hid_channel_size*size//2*size//2) -> (batch_size, latent_size)
        )

    def forward(self,x):
        x=self.layers(x)
        x=self.conversion_layer(x)
        return x

class layer_decoder(nn.Module):
    def __init__(self, size, base_channel_size, hid_channel_size, latent_size):
        super().__init__()
        self.size=size  # (batch_size, base_channel_size, size, size)
        self.base_channel_size=base_channel_size
        self.hid_channel_size=hid_channel_size
        self.latent_size=latent_size
        self.conversion_layer=nn.Sequential(
            nn.Linear(self.latent_size,self.hid_channel_size*self.size*self.size//4),
        )
        self.layers=nn.Sequential(
            nn.ConvTranspose2d(self.hid_channel_size, self.base_channel_size,kernel_size=3,padding=1, output_padding=1,stride=2), # (batch_size, base_channel_size, size, size) <- (batch_size, hid_channel_size, size//2, size//2)
            nn.ReLU(), # image values are between 0,1
        )

    def forward(self,x):
        x=self.conversion_layer(x)
        batch_size=x.shape[0]
        x=x.reshape(batch_size,self.hid_channel_size,self.size//2,self.size//2)
        x=self.layers(x)
        return x
        
class auto_encoder(nn.Module):
    def __init__(self,layer_encoder,layer_decoder, size, base_channel_size, hid_channel_size, latent_size):
        super().__init__()
        self.size=size
        self.base_channel_size=base_channel_size
        self.hid_channel_size=hid_channel_size
        self.latent_size=latent_size
        self.encoder=layer_encoder(size, base_channel_size, hid_channel_size, latent_size)
        self.decoder=layer_decoder(size, base_channel_size, hid_channel_size, latent_size)

    def forward(self,x):
            x=self.encoder(x)
            x=self.decoder(x)
            return x



auto_encoder=auto_encoder(layer_encoder,layer_decoder,size, base_channel_size, hid_channel_size, latent_size)
print(auto_encoder)

def initialize_weights(model):
    for name, pram in model.named_parameters():
        if "weight" in name:
            pram.data.normal_(std=torch.sqrt(torch.tensor(2)/(pram.shape[0]+pram.shape[1])))
        elif "bias" in name:
            pram.data.fill_(0)

initialize_weights(auto_encoder)

optimizer=optim.Adam(auto_encoder.parameters(), lr=lr)
loss_function=nn.MSELoss()

transform = transforms.Compose([
    transforms.ToTensor()
]) # convert the image to tensor

dataset= CIFAR10(root='./',train=True, download=True, transform =transform)
dataloader=data.DataLoader(dataset,batch_size=batch_size,shuffle=True)
sample_data, sample_label=next(iter(dataloader))
#print(sample_data) # (batch_size, base_channel_size, size, size)

class trainer(nn.Module):
    def __init__(self,model, dataloader, loss_function, optimizer, num_epoch=2):
        super().__init__()
        self.model=model
        self.dataloader=dataloader
        self.loss_function=loss_function
        self.optimizer=optimizer
        self.num_epoch=num_epoch
    
    def forward(self):
       model=self.model.to("mps")
       model.train()
       for epoch in torch.arange(self.num_epoch):
            epoch_loss=0
            for data, label in self.dataloader:
               data=data.to("mps")
               label=label.to("mps")
               self.optimizer.zero_grad()
               pred=model(data)
               loss=self.loss_function(pred, data)
               loss.backward()
               #nn.utils.clip_grad_norm_(model.parameters(),clip_grad_norm_val)
               self.optimizer.step()
               epoch_loss+=loss.item()
            print(f"Epoch {epoch+1} Loss: {epoch_loss/len(self.dataloader)}")

trainer=trainer(auto_encoder, dataloader, loss_function, optimizer, num_epoch)
trainer()