model_old.py

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
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):
        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()

    def forward(self, x):
        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.view(-1)        
        
        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.fc_out(x))
        # for i in range(14):
        #     x[i] = torch.exp(-x[i])

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

        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 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

        lr_max = 0.8       # learning rate
        lr_decay = 0.5 
        # lambda_L2 = 0.001
        lr_ind = np.arange(0,self.epochs+1)
        self.lr_vals = [lr_max/(1+lr_decay*lr_ind[i]) for i in lr_ind]
        # print(self.lr_vals)
        self.direc = '../data/1osc/tests/test1_'

        self.train_data, self.val_data = self.load_data(str('1osc/TRAIN_f'), self.N, self.val_size)
        self.train_size = len(self.train_data)
        self.val_size = len(self.val_data)
        self.cnn= CNN().to(device)
        # self.cnn = FF().to(device)
        self.loss_fn = nn.MSELoss(reduction='mean').to(device)
        self.optimiser = optim.Adadelta(self.cnn.parameters(), rho=0.7)
        # self.optimiser = optim.SGD(self.cnn.parameters(),lr=0.2)
        " 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 get_data(self,i):
        img, params = self.train_data[i]
        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 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 in range(self.train_size):
            # print("Train img: ", i)
            img, params = self.get_data(i)
            img_tens = torch.from_numpy(img)
            img_ = img_tens.float().unsqueeze(0).unsqueeze(0)
            self.optimiser.zero_grad()     # avoid accumulating gradients  
            y_train = torch.from_numpy(params).to(device)
            y_pred = self.cnn(img_.to(device))
            # y_pred = self.param_round(y_pred_).to(device)

            loss = self.loss_fn(y_train.unsqueeze(1), y_pred.unsqueeze(1))
            # loss = self.self.loss_fn_weighted(y_train,y_pred)
            loss.backward()           # calculate gradients
            self.optimiser.step()    
            loss_train_total += loss.item()
            # loss_train_total += loss
        
        return loss_train_total


    def val(self):
        loss_val_total  = 0

        for i in range(self.val_size):
            # print("Test img: ", i)
            img, params = self.get_data(i)
            img_tens = torch.from_numpy(img)
            img_ = img_tens.float().unsqueeze(0).unsqueeze(0)
            y_val = torch.from_numpy(params).to(device)
            y_pred = self.cnn(img_.to(device))

            val_loss = self.loss_fn(y_val.to(device), y_pred)
            # val_loss = self.self.loss_fn_weighted(y_val,y_pred)
            loss_val_total += val_loss.item()
            # loss_val_total += val_loss

        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.optimiser = optim.SGD(self.cnn.parameters(), lr=self.lr_vals[epoch])

            self.cnn.train()             # set the self.cnn to training mode 
            loss_train_total = self.train_step()

            # L2 = 0
            # for p in self.cnn.parameters():
            #     L2 += torch.sum(torch.mul(p,p))

            # L2_tot = L2*lambda_L2
            
            losses_train.append(loss_train_total/self.train_size)
            
            # reglosses_train.append((loss_train_total+L2_tot)/len(train_data))

            with torch.no_grad():   
                self.cnn.eval()
                loss_val_total = self.val()
                losses_val.append(loss_val_total/self.val_size)
                # reglosses_val.append((loss_val_total+L2_tot)/len(val_data))

            # if epoch % 10 == 0 and epoch > 0: 
            #     print('save the model at epoch {}'.format(epoch))
            #     # torch.save(self.cnn, "../data/weights/weightsSize{}Epoc{}.pth".format(N,epoch))
            #     torch.save(self.cnn, "../data/"+source+"_weights/"+output+"/weights{}.pth".format(epoch))

        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))
        
        # torch.save(reglosses_train, '../data/'+source+'/tests/'+output+'reglosses_train_{}'.format(epoch))
        # torch.save(reglosses_val, '../data/'+source+'/tests/'+output+'reglosses_val_{}'.format(epoch))
        
        return losses_train, losses_val
        

    def load_data(self,folder, N, val_size):
        direc_ = str('../data/' + folder + '/')
        data_list_train = os.listdir(direc_)[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(direc_ + data_list_train[i], allow_pickle=True)
            train_data.append(img)

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

        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,133]) 
    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()
    print('Total Test Loss = ', test_loss)