model_batch2.py

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torch.utils.data import DataLoader
import numpy as np
import os
import matplotlib.pyplot as plt
import random
import time
import cv2 as cv

device = 'cuda:0' if torch.cuda.is_available() else 'cpu'  # check whether a GPU is available

class CNN(nn.Module):
    
    def __init__(self, num_channels):
        super(CNN, self).__init__()
        knum = 64
        self.conv1 = nn.Conv2d(1, knum, 5, padding=2, bias=True)
        self.conv2 = nn.Conv2d(knum,knum, 5, padding=2, bias=True)
        self.conv3 = nn.Conv2d(knum, 1, 5, padding=2, bias=True)

        self.pool = nn.MaxPool2d(2,2)
        self.insize = 5000
        self.hsize1 = 1000
        self.hsize2 = 200
        # self.hsize3 = 40
        self.hsize4 = 50
        self.outsize = 14

        self.fc_in = nn.Linear(self.insize, self.hsize1) 
        self.fc1 = nn.Linear(self.hsize1, self.hsize2)
        self.fc2 = nn.Linear(self.hsize2, self.hsize4)
        # self.fc3 = nn.Linear(self.hsize3, self.hsize4)
        self.fc_out = nn.Linear(self.hsize4, self.outsize)

        self.softmax = nn.Softmax(dim=0)
        self.relu = nn.ReLU()
        self.logsig = nn.LogSigmoid()
        self.sigmoid = nn.Sigmoid()
        self.dropout = nn.Dropout(0.5)


    def forward(self, x):
        # print(x.size())
        x = x.type(torch.cuda.FloatTensor)
        x = self.relu(self.conv1(x))
        x = self.pool(x)    

        x = self.relu(self.conv2(x))
        x = self.pool(x)
   
        x = self.relu(self.conv3(x))
        x = self.pool(x)

        # print('SIZE',x.size())

        x = x.flatten()      

        x = self.fc_in(x)
        # x = self.dropout(x)
        x = self.relu(x)
        x = self.fc1(x)
        # x = self.dropout(x)
        x = self.relu(x)

        x = self.fc2(x)
        # x = self.dropout(x)
        x = self.relu(x)
        # print(x.size())

        x = self.fc_out(x)  
        x = self.relu(x)
        
        x = 0.07*x
        x = self.sigmoid(x)
        # a = 4.32790682748   # solve for y=1 when x=1
        a = 2
        x = a*(x-0.5)
        # print(x)
        # torch.clamp(x,0,1)

        return x


class FF(nn.Module):
    
    def __init__(self):
        super(FF, self).__init__()
        self.insize = 335067
        self.hsize1 = 1000
        self.hsize2 = 500
        self.hsize3 = 100
        self.hsize4 = 50
        self.hsize5 = 200
        self.hsize6 = 50
        self.outsize = 14

        self.fc_in = nn.Linear(self.insize, self.hsize1)
        self.fc1 = nn.Linear(self.hsize1,self.hsize2)
        self.fc2 = nn.Linear(self.hsize2,self.hsize3)
        self.fc3 = nn.Linear(self.hsize3,self.hsize4)
        # self.fc4 = nn.Linear(self.hsize4,self.hsize5)
        # self.fc5 = nn.Linear(self.hsize5,self.hsize6)
        self.fc_out = nn.Linear(self.hsize4,self.outsize)

        # self.fc3 = nn.Linear(self.hsize2,self.outsize)
        self.relu = nn.ReLU()
        self.sigmoid = nn.Sigmoid()
        self.softmax = nn.Softmax(dim=0)

        
    def forward(self, x):
        x = torch.flatten(x)
        x = self.relu(self.fc_in(x))
        x = self.relu(self.fc1(x))
        x = self.relu(self.fc2(x))
        x = self.relu(self.fc3(x))
        # x = self.relu(self.fc4(x))
        # x = self.relu(self.fc5(x))
        x = self.relu(self.fc_out(x))

        x = 0.07*x
        x = self.sigmoid(x)
        x = 2*(x-0.5)

        return x


class CreateDataset(DataLoader):
    """Face Landmarks dataset."""

    def __init__(self, file_dir, data_list, N):
        """
        Args:
            csv_file (string): Path to the csv file with annotations.
            root_dir (string): Directory with all the images.
            transform (callable, optional): Optional transform to be applied
                on a sample.
        """
        self.file_dir = file_dir
        self.data_list = data_list
        self.N = N

    def __len__(self):
        return len(self.data_list)

    def __getitem__(self, idx):
        print(idx)
        img, params = get_data(self.file_dir,idx)
        return img, params


def get_data(file_dir,i):
    img, params = np.load(file_dir + 'img{}.npy'.format(i), allow_pickle=True)
    params = params/127
    img_copy = np.zeros([5001,67]) 
    img_norm = cv.normalize(img, img_copy,0,255,cv.NORM_MINMAX)

    # img_norm = img / np.linalg.norm(img)
    # img_norm = img/(np.max(img)-np.min(img))
    # img_norm = img/sum(img)

    return img_norm, params


def weights_init_uniform(m):
    classname = m.__class__.__name__
    # for every Linear layer in a model
    if classname.find('Conv') != -1:
        # apply a uniform distribution to the weights and a bias=0
        # print('BEFORE',m.weight.data.size())
        m.weight.data.uniform_(-0.5,0.5)
        m.bias.data.fill_(0)

    if classname.find('Linear') != -1:
        # apply a uniform distribution to the weights and a bias=0
        m.weight.data.uniform_(-0.1, 0.1)
        m.bias.data.fill_(0)


class Train():
    def __init__(self):
        self.epochs = 10        # number of self.epochs
        self.val_size = 0.2     # proportion of train data to use for validation
        self.N = 100           # size of data set
        self.bsize = 1
        self.nworkers = 1

        "LOAD DATA"
        self.direc = '../data/1osc/tests/test0_'    # save direc
        self.root = 'D:/Documents/UNI/Year 4/Summer Project/data/1osc/TRAIN_f/'

        print('Loading Data...')
        self.trainLoader, self.valLoader = self.load_data(self.root, self.N, self.val_size)

        # self.train_list, self.val_list = self.load_data(str(self.root+'TRAIN_f'), self.N, self.val_size)
        # self.train_data = CreateDataset(str(self.root+'TRAIN_f/'),self.train_list, self.N)
        # self.val_data = CreateDataset(str(self.root+'TRAIN_f/'),self.val_list, self.N)
        # self.trainLoader=torch.utils.data.DataLoader(self.train_data, batch_size=self.bsize, shuffle=True, num_workers=self.nworkers)
        # self.valLoader=torch.utils.data.DataLoader(self.val_data, batch_size=self.bsize, shuffle=True, num_workers=self.nworkers)
        print('Finished')
        self.train_size = len(self.trainLoader)
        self.val_size = len(self.valLoader)
        print(self.train_size)
        print(self.val_size)
        self.cnn= CNN(self.bsize).to(device)
        # self.cnn = FF().to(device)

        "LOSS FN"
        self.loss_fn = nn.MSELoss(reduction='mean').to(device)
        # self.loss_fn = nn.CrossEntropyLoss()

        self.optimiser = optim.Adadelta(self.cnn.parameters(), rho=0.7,lr=1)
        # self.optimiser = optim.SGD(self.cnn.parameters(),lr=0.2)
        # self.optimiser = optim.RMSprop(self.cnn.parameters(),lr=0.01)

        "INITIALISE WEIGHTS"
        # lay = self.cnn.conv2.weight.data
        # print(torch.min(lay),torch.max(lay))
        # self.cnn.apply(weights_init_uniform)
        # lay = self.cnn.conv2.weight.data
        # print(torch.min(lay),torch.max(lay))

        " define loss weights "

        # self.weights = np.ones(14)
        # self.weights[1]=10    # waveform  
        # self.weights[0]=10    # range
        # self.weights[-1]=5   # tuning
        # self.weights[5]=2    # filt A
        # self.weights[6]=2    # filt D
        # self.weights[9]=2    # amp A 
        # self.weights[10]=2  # amp D


        self.main_train()


    def param_round(self,params):
        param_list = [0,1,14,15,16,17,18,19,20,21,22,23,24,25]
        param_six = [0,1,3,5,7,9]           # ind for six-pos params
        param_two = [19,23]                 # ind for binary params

        vals_six = [0,26,51,77,102,127]     # vals for six-pos params
        vals_two = [0,127]                  # vals for binary params

        params_out = params.clone()

        for j in range(14):
            i = param_list[j]
            r = params[j]

            if i in param_six:
                val_ind = int(torch.round(5*r))
                val = vals_six[val_ind]

            elif i in param_two:
                val_ind = int(torch.round(r))
                val = vals_two[val_ind]

            else:
                val = int(torch.round(127*r))

            params_out[j] = val/127
        return params_out


    def loss_fn_weighted(self, x, y):    
        loss=0

        for i in range(14):
            loss += weights[i]*(x[i]-y[i])**2
        
        return loss/14

 
    def train_step(self):
        loss_train_total = 0
        
        for i, sample in enumerate(self.trainLoader):
            # print("Train img: ", i)
            # img_b, params_b = sample
            # for j in range(self.bsize):

            # img = img_b[j]
            img, params = get_data(self.root,i)
            img_tens = torch.from_numpy(img)
            img_ = img_tens.float().unsqueeze(0).unsqueeze(0)

            # params = params_b[j]
            self.optimiser.zero_grad()     # avoid accumulating gradients  
            y_train = torch.from_numpy(params).to(device)

            # print(img_.size())
            y_pred = self.cnn(img_.to(device))
            # y_pred = self.param_round(y_pred_)

            loss = self.loss_fn(y_train, y_pred)

            # loss = loss/self.bsize
            loss.backward()           # calculate gradients
            self.optimiser.step()    

            loss_train_total += loss.item()

        return loss_train_total


    def val(self):
        loss_val_total  = 0 
        for i, sample in enumerate(self.valLoader):
            # img_b, params_b = sample
            # for j in range(self.bsize):
                # img = img_b[j]
            img, params = get_data(self.root,i)

            img_tens = torch.from_numpy(img)

            # params = params_b[j]
            y_val = torch.from_numpy(params).to(device)
            img_ = img_tens.float().unsqueeze(0).unsqueeze(0)

            y_pred = self.cnn(img_.to(device))
            # y_pred = self.param_round(y_pred_)

            loss = self.loss_fn(y_val, y_pred)

            # val_loss = self.self.loss_fn_weighted(y_val,y_pred)
            # loss = loss/self.bsize
            loss_val_total += loss.item()

        return loss_val_total


    def main_train(self):
        losses_train = list()
        losses_val   = list()
        # reglosses_train = list()
        # reglosses_val = lis0t()

        for epoch in range(self.epochs+1):
            print('Epoch: ', epoch)
            self.cnn.train()             # set the self.cnn to training mode 
            loss_train_total = self.train_step()
            
            losses_train.append(loss_train_total/self.train_size)
            
            with torch.no_grad():   
                self.cnn.eval()
                loss_val_total = self.val()
                losses_val.append(loss_val_total/self.val_size)

        torch.save(self.cnn, self.direc+'weights{}.pth'.format(epoch))
        print('SAVING', self.direc)
        torch.save(losses_train, self.direc+'losses_train_{}'.format(epoch))
        torch.save(losses_val, self.direc+'losses_val_{}'.format(epoch))
        
        return losses_train, losses_val
        

    def load_data(self,folder, N, val_size):
        data_list_train = os.listdir(folder)[0:N]

        validation = random.sample(data_list_train,k=int(self.val_size*N))
        for item in validation:
            data_list_train.remove(item)

        data_list_val = validation
        
        train_data = []
        val_data = []
        
        for i in range(len(data_list_train)):
            img = np.load(folder + data_list_train[i], allow_pickle=True)
            train_data.append(img)

        for i in range(len(data_list_val)):
            img = np.load(folder + data_list_val[i], allow_pickle=True)
            val_data.append(img)

        # return data_list_train, data_list_val
        return train_data, val_data

    def test_all(self):
        torch.cuda.empty_cache()
        self.cnn.eval()
        total_loss = 0
        with torch.no_grad():
            for i in range(100):
                img, params = load_img(i)
                output = self.cnn(img.to(device)).to('cpu')
                # print(params)
                # print(output)
                loss = self.loss_fn(torch.tensor(params), output)
                total_loss += loss

        return total_loss

def load_img(n):
    img, params = np.load('../data/1osc/testdata/TRAIN_f/' + 'img{}.npy'.format(int(n)), allow_pickle=True)
    # plt.imshow(img)
    # plt.show()
    params = params/127
    img_copy = np.zeros([5001,67]) 
    img_norm = cv.normalize(img, img_copy,0,255,cv.NORM_MINMAX)

    img_tens = torch.from_numpy(img_norm)
    img_ = img_tens.float().unsqueeze(0).unsqueeze(0)

    return img_, params

if __name__ == '__main__':
    torch.cuda.empty_cache()

    start_time = time.time()
    T = Train()
    elapsed_time = time.time() - start_time
    print("Training Time:",elapsed_time/60," Min")

    # test_loss = T.test_all()/100
    # print('Total Test Loss = ', test_loss)