chouse_train/train_model_1.2.py

import torch
import torchvision.models
import sys
import pandas as pd
import os
from datasetmaker import TrainDataClass, TestDataClass
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms, datasets
from torch import optim
import torch.nn as nn
import torch.nn.functional as F

weight_dir_root = 'weights_backup'
acc_dir_root = 'accs'


'''
filename:train_model_1.0.py
create date: 11/17/2020 Tue
Author: Jeong Geol Kim
Contact: atarib4816@gmail.com
Description: This .py file will train, save weight and save Results with parameters. parameters will input by .txt file.
filename format: partofarchitecture_backbone_yymmdd.lcfg
eg) shapesort_mobilenetv2_201118.txt
'''


def Read_Config(lcfg_name):
    '''
    Description: read options from .lcfg and initialize learning options.
    .lcfg file has 10 pramas: {role} {bbname} {dsversion} {dsname} {epoch} {batchsize} {optmizer} {learningloss} {lossfunction} {saveperiod}
    '''       
    N_OF_PARAMS = 10
    dict_key = ['role', 'bbname', 'dsversion', 'dsname', 'epoch', 'batchsize', 'optimizer', 'learningloss', 'lossfunction', 'saveperiod']

    f = open(lcfg_name, 'r')
    p_string = f.readline()
    f.close()
    params = p_string.split()

    if len(params) != N_OF_PARAMS: 
        print("Error: {} is broken. please check learning configures. Or you have to change constant N_OF_PARMAS.".format(lcfg_name))
        return
    else:
        print("----Learning Options----")
        print("Role of Classifier: {}".format(params[0]))
        print("Backbone Name: {}".format(params[1]))
        print("Dataset version: {}".format(params[2]))
        print("Dataset target: {}".format(params[3]))
        print("Number of iterations: {}".format(params[4]))
        print("Images per Batch: {}".format(params[5]))
        print("Optimizer Name: {}".format(params[6]))
        print("learning loss: {}".format(params[7]))
        print("Loss function Name: {}".format(params[8]))
        print("Period of saving weight: {}".format(params[9]))
        parmas_dictionary = dict(zip(dict_key, params))
        
        return parmas_dictionary


def Preprocess(params, lcfg_name):
    dataset_root = os.path.join(os.getcwd(),'datasets',params['dsversion'],params['dsname'])
    classes = tuple(os.listdir(os.path.join(dataset_root,'train')))

    #NOTE: we don't have to transform because we already shrink raw images as same as MoblieNetV2's input size
    tsfm=transforms.Compose([
        transforms.ToTensor()
    ])

    train = TrainDataClass(dataset_root, tsfm)
    trainloader = DataLoader(train, batch_size=int(params['batchsize']), shuffle=True)

    test = TestDataClass(dataset_root, tsfm)
    testloader = DataLoader(test, batch_size=int(params['batchsize']), shuffle=True)

    return classes, trainloader, testloader




def Model_Construct(params, lcfg_name):
    if params['bbname'] == 'moblienetv2':
        model = torch.hub.load('pytorch/vision', 'mobilenet_v2', pretrained=True)
        return model
    if params['bbname'] == 'resnet50':
        model = torch.hub.load('pytorch/vision', 'resnet50', pretrained=True)
        return model
    

def Train(params, lcfg_name, model, trainLoader, testLoader, criterion, optimizer):
    params = params
    lcfg_name = lcfg_name
    model = model
    trainLoader = trainLoader
    criterion = criterion
    optimizer = optimizer
    weight_dir_root = 'weights_backup'
    acc_dir_root = 'accs'

    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    #print(torch.cuda.device_count())
    print(device)
    model = model.to(device)
                
    for epoch in range(int(params['epoch'])):    
        running_loss = 0.0
        for i, data in enumerate(trainLoader, 0):
            # [inputs, labels]의 목록인 data로부터 입력을 받은 후;
            inputs, labels = data[0].to(device), data[1].to(device)

            # 변화도(Gradient) 매개변수를 0으로 만들고
            optimizer.zero_grad()

            # 순전파 + 역전파 + 최적화를 한 후
            outputs = model(inputs).to(device)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()

            # 통계를 출력합니다.
            running_loss += loss.item()
            #print(i)
            if i % 50 == 49:    # print every 50 mini-batches
                print('[%d, %5d] loss: %.3f' %
                      (epoch + 1, i + 1, running_loss / 2000))                
                running_loss = 0.0
        correct = 0
        total = 0
        with torch.no_grad():
            for data in testLoader:
                inputs, labels = data[0].to(device), data[1].to(device)
                outputs = model(inputs).to(device)
                _, predicted = torch.max(outputs.data, 1)
                total += labels.size(0)
                correct += (predicted == labels).sum().item()
        accuracy = (100 * correct / total)
        print('Accuracy: {0}'.format(accuracy))

        acc_name = os.path.join(os.getcwd(),acc_dir_root,lcfg_name + '.csv')
        weight_name = os.path.join(os.getcwd(), weight_dir_root,lcfg_name,params['role']+'_'+str(epoch+1)+'.pth')

        if epoch % int(params['saveperiod']) == (int(params['saveperiod']) - 1):
            params['savedat'] = epoch + 1
            params['accuracy'] = accuracy
            cols = params.keys()
            #NOTE: dictionary needs extend
            if not os.path.isfile(acc_name):                
                df = pd.DataFrame(index = range(0), columns = cols)
                row = list(params.values())
                
                df = df.append(pd.Series(row, index=cols), ignore_index= True)
                df.to_csv(acc_name)
            else:
                df = pd.read_csv(acc_name, index_col = 0)                
                row = list(params.values())                
                df = df.append(pd.Series(row, index=cols), ignore_index= True)
                df.to_csv(acc_name)

            if not os.path.isfile(weight_name):
                torch.save(model.state_dict(),weight_name)
        


        #if epoch % int(params['saveperiod']) == (int(params['saveperiod']) - 1):            
        #    torch.save(model.state_dict(), os.path.join(os.getcwd(), weight_dir_root,lcfg_name))

    print('Finished Training')
    return model


'''
def Test(model, testLoader):
    correct = 0
    total = 0
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

    with torch.no_grad():
        for data in testLoader:
            inputs, labels = data[0].to(device), data[1].to(device)
            outputs = model(inputs).to(device)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
    accuracy = (100 * correct / total)
    print('Accuracy: {}'.format(accuracy))
    return accuracy
'''



def Save_Weight(params, lcfg_name):
    DEFAULT_DIR = 'weights_backup'
    pass


def Save_Result(params, lcfg_name):
    DEFAULT_DIR = 'acc_save'
    result_file = os.path.join(os.getcwd(),DEFAULT_DIR,lcfg_name) + '.csv'
    
    if os.path.isfile(result_file):
        df = pd.read_csv(result_file)
    else:
        cols = ['backbone', 'dataset', 'epoch', 'batchsize', 'optimizer', 'learningloss', 'lossfunction', 'saveat', 'accuracy', 'precision', 'recall']
        #row = []
        df = pd.DataFrame(params, columns = cols)
        print(df)
        df.to_csv(result_file)

    pass

def Set_Optimizer(params, model):    
    if params['optimizer'] == 'RMSprop':
        optimizer = optim.RMSprop(model.parameters(), lr = float(params['learningloss']), momentum=0.9)
        return optimizer
    elif params['optimizer'] == 'SGD':
        optimizer = optim.SGD(model.parameters(), lr = float(params['learningloss']), momentum=0.9)
        return optimizer
    elif params['optimizer'] == 'Adam':
        optimizer = optim.Adam(model.parameters(), lr = float(params['learningloss']))
        return optimizer
        
        
    




if __name__ == "__main__":    
    for root, dirs, files in os.walk('.', topdown = True):
        for filename in files:
            if filename.endswith(('lcfg')):                                
                lcfg_path = os.path.join(root, filename)
                lcfg_name = os.path.splitext(filename)[0]
                
                params = Read_Config(lcfg_path)
                #print(params)
                if not os.path.exists(os.path.join(os.getcwd(),weight_dir_root,lcfg_name)):
                    os.makedirs(os.path.join(os.getcwd(),weight_dir_root,lcfg_name), exist_ok=True)
                else:
                    pass

                classes, trainLoader, testLoader = Preprocess(params, lcfg_name)
                #print(classes)

                model = Model_Construct(params, lcfg_name)

                #print(model.eval())
                if params['lossfunction'] == 'CrossEntropyLoss':
                    criterion = nn.CrossEntropyLoss()
                else:
                    criterion = nn.MultiLabelMarginLoss()

                optimizer = Set_Optimizer(params, model)
                
                model = Train(params, lcfg_name, model, trainLoader, testLoader, criterion, optimizer)               
            else:
                pass