train_Temporal_3ch.py

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision.transforms as transforms
import torch.nn.functional as F
import os
import pandas as pd
import matplotlib.pyplot as plt
from torch.utils.data import Dataset                  
from torchvision import datasets, transforms
from Dataloder.MotionHistory_dataloder import MotionDataset
from Dataloder.Pseudo_MotionHistory_dataloder import PseudoMotionHistory_Dataset
from torch.autograd import Variable
from PIL import Image
from statistics import mean
from torch.utils.data.dataset import Subset
from Network.Temporal.MobileNet import MobileNet_V2_Temporal
from Network.Temporal.VGG16 import VGG16_Temporal
from sklearn.model_selection import train_test_split

EpochNum = 51                                          
Height = 224                                            
Width = 224                                           
BatchSize = 16   
LearningRate = 0.001                    
TrainPath = r"" 
TestPath =r""
modelPath = r""          
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")                         


class Trainer():
    def __init__(self,model, optimizer, criterion, trainLoader, valLoader, transform):
        self.model = model
        self.optimizer = optimizer
        self.criterion = criterion
        self.trainLoader = trainLoader
        self.valLoader = valLoader
        self.transform = transform
        self.totalTrainLoss = []
        self.TrainCorrect = []
        self.totalValLoss = []
        self.ValCorrect = []
        self.max_acc = 0
        self.Fig = plt.figure(figsize=[10,10])
    def Train(self,epoch):
        self.model.train()
        train_loss,train_acc = 0.0,0.0
        t_loss,t_acc =  0.0,0.0
        train_log = ""
        for batchIdx,(img,label) in enumerate(self.trainLoader):
            img,label = Variable(img.cuda()),Variable(label.cuda())
            output = self.model(img)
            loss = self.criterion(output,label)
            train_loss += loss.data.item()
            self.optimizer.zero_grad()
            loss.backward()
            self.optimizer.step()
            pred = output.data.max(dim=1)[1]
            train_acc += pred.eq(label.data).cpu().sum()
            t_loss = train_loss/((batchIdx+1)*BatchSize)
            t_acc  = 100*train_acc.data.item() / ((batchIdx+1)*BatchSize)
            train_log = "epoch : {:3} train_loss : {:3.10} train_acc : {:3.10}".format(str(epoch+1), str(t_loss), str(t_acc))
            print("\r"+train_log,end="")
        self.totalTrainLoss.append(t_loss)
        self.TrainCorrect.append(t_acc)

        self.model.eval()
        val_loss,val_acc,min_acc = 0.0,0.0,0.0
        v_loss,v_acc = 0.0,0.0
        val_log = ""
        with torch.no_grad():
            for batchIdx,(img,label) in enumerate(self.valLoader):
                img,label = Variable(img.cuda()),Variable(label.cuda())
                output = self.model(img)
                loss = self.criterion(output,label)
                val_loss += loss.data.item()
                pred = output.data.max(dim=1)[1]
                val_acc += pred.eq(label.data).cpu().sum()
                v_loss = val_loss/((batchIdx+1)*BatchSize)
                v_acc = 100*val_acc.data.item() / ((batchIdx+1)*BatchSize)
                val_log = train_log + " val_loss : {:3.10} val_acc : {:3.10}".format(str(v_loss), str(v_acc))
                print("\r"+val_log,end="")
            self.totalValLoss.append(v_loss)
            self.ValCorrect.append(v_acc)
        print()
        if v_acc > self.max_acc:
            self.max_acc = v_acc
            with open("models/"+str(modelPath)+"/"+str(modelPath)+".pth", "wb") as savePath:
                torch.save({'epoch': epoch,
                            'model_state_dict': model.state_dict(),
                            'optimizer_state_dict': optimizer.state_dict(),
                            'loss': loss,},savePath)			
    def graph_plot(self):
        plt.clf()
        plt.style.use('ggplot')
        lossFig = self.Fig.add_subplot(2,1,1) 
        plt.title('Loss Graph')
        plt.plot(self.totalTrainLoss, label='train loss')
        plt.plot(self.totalValLoss, label='validation loss')
        plt.xlabel('Epoch')
        plt.ylabel('Loss')
        plt.legend()
        accuracyFig = self.Fig.add_subplot(2,1,2) 
        plt.title('Accuracy Graph')
        plt.minorticks_on()
        plt.ylim(0,100)
        plt.plot(self.TrainCorrect, label='train acc')
        plt.plot(self.ValCorrect, label='validation acc')
        plt.xlabel('Epoch')
        plt.ylabel('Accuracy')
        plt.legend()

    def main(self):
        for epoch in range(EpochNum):
            self.Train(epoch)
            self.graph_plot()
            self.Fig.savefig("models/"+str(modelPath)+"/"+str(modelPath)+".png")

if __name__ == '__main__':
    transform = transforms.Compose([
        transforms.Resize([Height, Width]),
        transforms.ToTensor(),
    ])
    os.mkdir("models/"+str(modelPath))
    criterion = nn.CrossEntropyLoss()
    Traindataset = PseudoMotionHistory_Dataset(TrainPath, transform, 5)
    Testdataset = PseudoMotionHistory_Dataset(TrainPath, transform, 5)
    print("Load TrainDataSet 「"+str(TrainPath)+"」")
    print("Load TestDataSet 「"+str(TestPath)+"」")
    torch.backends.cudnn.benchmark = True
    MobileNet = MobileNet_V2_Temporal()
    model = MobileNet.model.cuda()
    model.to(device)
    optimizer = optim.Adam(model.parameters(),lr=0.001)
    print(len(Traindataset),len(Testdataset))
    trainLoader = torch.utils.data.DataLoader(Traindataset, batch_size=BatchSize, shuffle=True, num_workers=os.cpu_count(), pin_memory=True)
    valLoader = torch.utils.data.DataLoader(Testdataset, batch_size=BatchSize, shuffle=False, num_workers=os.cpu_count(), pin_memory=True)
    train = Trainer(model, optimizer, criterion,trainLoader, valLoader, transform)
    train.main()