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