Cifar10 Autoencoder 20230710 Rayleigh

import math
import torch
import random
import torchvision
import torch.nn as nn
from tqdm import tqdm
import torch.optim as optim
import torch.nn.functional as f
import matplotlib.pyplot as plt
import torchvision.transforms as tr
from torch.utils.data import DataLoader
import numpy as np

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

print(device)

args = {
    'BATCH_SIZE' : 50,
    'LEARNING_RATE' : 0.001,
    'NUM_EPOCH' : 20,
    'SNR_dB' : [1, 10, 20],
    'latent_dim' : 500,
    'input_dim' : 32 * 32
}

transf = tr.Compose([tr.ToTensor(), tr.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

trainset = torchvision.datasets.CIFAR10(root = './data', train = True, download = True, transform = transf)
testset = torchvision.datasets.CIFAR10(root = './data', train = False, download = True, transform = transf)

trainloader = DataLoader(trainset, batch_size = args['BATCH_SIZE'], shuffle = True)
testloader  = DataLoader(testset, batch_size = args['BATCH_SIZE'], shuffle = False)

def transmit_img_AE(model_name, SNRdB, testloader) :

    model = Autoencoder()
    model.load_state_dict(torch.load(model_name))

    fig = plt.figure()

    batch_size = args['BATCH_SIZE']

    for data in testloader :

        inputs = data[0]

        outputs = model(inputs, SNRdB = SNRdB)

        break

    rand_num = random.randrange(batch_size)

    plt.subplot(1, 2, 1)
    plt.imshow( inputs[rand_num].permute(1, 2, 0).detach().numpy())
    plt.title('Original img')

    plt.subplot(1, 2, 2)
    plt.imshow(outputs[rand_num].permute(1, 2, 0).detach().numpy())
    plt.title('Decoded img')




def PSNR_test(model_name, testloader) :

    if model_name[-6] == '=' :
        SNRdB = int(model_name[-5])
    else :
        SNRdB = int(model_name[-6] + model_name[-5])

    model_1 = Autoencoder().to(device)
    model_1.load_state_dict(torch.load(model_name))

    PSNR_list = []

    for data in testloader :
        inputs = data[0].to(device)

        outputs = model_1( inputs, SNRdB = SNRdB)


        for j in range(len(inputs)) :
            PSNR = 0
            for k in range(3) :

                MSE = torch.sum(torch.pow(inputs[j][k] - outputs[j][k], 2)) / args['input_dim']
                PSNR += 10 * math.log10(1 / MSE)

            PSNR_list.append(PSNR / 3)
    print("PSNR : {}".format(round(sum(PSNR_list) / len(PSNR_list), 2)))

    return round(sum(PSNR_list) / len(PSNR_list), 2)



def compare_PSNR(SNRdB_list, model_name_without_SNRdB, testloader) :

    PSNR_list = []

    for i in range(len(SNRdB_list)) :
        SNRdB = SNRdB_list[i]
        model_name = model_name_without_SNRdB + "_SNR=" + str(SNRdB) + ".pth"
        PSNR = PSNR_test(model_name, testloader)
        PSNR_list.append(PSNR)

    plt.plot(args['SNR_dB'], PSNR_list, linestyle = 'dashed', color = 'blue', label = "Rayleigh")
    plt.grid(True)
    plt.legend()
    plt.show()



def model_train(SNRdB_list, learning_rate, epoch_num, trainloader) :

    for i in range( len(args['SNR_dB']) ):

        model = Autoencoder().to(device)
        criterion = nn.MSELoss()
        optimizer = optim.Adam(model.parameters(), lr = learning_rate)

        print("+++++ SNR = {} Training Start! +++++".format(args['SNR_dB'][i]))

        for epoch in range(epoch_num) :

            running_loss = 0.0

            for data in trainloader :
                
                inputs = data[0].to(device)

                optimizer.zero_grad()
                outputs = model( inputs, SNRdB = args['SNR_dB'][i])
                loss = criterion(inputs, outputs)
                loss.backward()
                optimizer.step()
                running_loss += loss.item()

            cost = running_loss/len(trainloader)

            print("[{} Epoch] PSNR : {}".format(epoch + 1, round(10 * math.log10(1 / cost), 3)))
        print()
        PATH = "./"
        torch.save(model.state_dict(), PATH + "model_AWGN(color)" + "_SNR=" + str(args['SNR_dB'][i]) +".pth")
        

class Encoder(nn.Module):
    def __init__(self):
        super(Encoder, self).__init__()
        c_hid = 32

        self.encoder = nn.Sequential(
            nn.Conv2d(3, c_hid, kernel_size=3, padding=1, stride=2),  # 32x32 => 16x16
            nn.ReLU(),
            nn.Conv2d(c_hid, c_hid, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.Conv2d(c_hid, 2 * c_hid, kernel_size=3, padding=1, stride=2),  # 16x16 => 8x8
            nn.ReLU(),
            nn.Conv2d(2 * c_hid, 2 * c_hid, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.Conv2d(2 * c_hid, 2 * c_hid, kernel_size=3, padding=1, stride=2),  # 8x8 => 4x4
            nn.ReLU(),
            nn.Flatten(),  # Image grid to single feature vector
            nn.Linear(2 * 16 * c_hid, args['latent_dim'])
        )

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




class Decoder(nn.Module):
    def __init__(self):
        super(Decoder, self).__init__()
        c_hid = 32

        self.linear = nn.Sequential(
            nn.Linear(args['latent_dim'], 2 * 16 * c_hid),
            nn.ReLU()
            )

        self.decoder = nn.Sequential(
            nn.ConvTranspose2d(2 * c_hid, 2 * c_hid, kernel_size=3, output_padding=1, padding=1, stride=2),  # 4x4 => 8x8
            nn.ReLU(),
            nn.Conv2d(2 * c_hid, 2 * c_hid, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.ConvTranspose2d(2 * c_hid, c_hid, kernel_size=3, output_padding=1, padding=1, stride=2),  # 8x8 => 16x16
            nn.ReLU(),
            nn.Conv2d(c_hid, c_hid, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.ConvTranspose2d(c_hid, 3, kernel_size=3, output_padding=1, padding=1, stride=2),  # 16x16 => 32x32
        )

    def forward(self, x):

        x = self.linear(x)
        x = x.reshape(x.shape[0], -1, 4, 4)
        decoded = self.decoder(x)

        return decoded



class Autoencoder(nn.Module):
    def __init__(
            self,
            encoder_class: object = Encoder,
            decoder_class: object = Decoder
    ):
        super(Autoencoder, self).__init__()

        self.encoder = encoder_class()
        self.decoder = decoder_class()


    def Rayleigh(self, input, SNRdB):

        normalized_tensor = f.normalize(input, dim=1)

        SNR = 10.0 ** (SNRdB / 10.0)
        K = args['latent_dim']
        std = 1 / math.sqrt(K * SNR)

        h = (np.sqrt(torch.normal(0, 1, size=normalized_tensor.size())**2 + torch.normal(0, 1, size=normalized_tensor.size())**2) / np.sqrt(2)).to(device)
        n = torch.normal(0, std, size=normalized_tensor.size()).to(device)

        return normalized_tensor + n/h


    def forward(self, x, SNRdB):

        encoded = self.encoder(x)
        #print("encoded : {}".format(encoded[0]))
        #print("encoded size : {}".format(encoded.size()))

        encoded_AWGN = self.Rayleigh(encoded, SNRdB)
        #print("encoded_AWGN size : {}".format(encoded_AWGN.size()))        

        decoded = self.decoder(encoded_AWGN)
        #print("decoded : {}".format(decoded[0]))
        #print("decoded size : {}".format(decoded.size()))

        return decoded
        


SNRdB_list = args['SNR_dB']
learning_rate = args['LEARNING_RATE']
epoch_num = args['NUM_EPOCH']
trainloader = trainloader
model_name_without_SNRdB = 'model_AWGN(color)'

model_train(SNRdB_list, learning_rate, epoch_num, trainloader)

compare_PSNR(SNRdB_list, model_name_without_SNRdB, testloader)