train.py

import torch 
import numpy as np
import random
from model.image_cnn import CNN
from trainer import * 
# from utils import config
import utils
from sklearn import metrics
import torch.nn as nn
import os
from matplotlib import pyplot as plt
from IPython.display import clear_output
import importlib
from sklearn.model_selection import ParameterGrid
import collections
import yaml 
from experiment import Experiment 
# from apex import amp, fp16_utils
import re 
import argparse
parser = argparse.ArgumentParser(description='')
parser.add_argument('--model_name', type=str, required=True)
parser.add_argument('--model_type', type=str, required=True)
parser.add_argument('--gpu', type=str, required=True, default = "0,1")
parser.add_argument('--budget', type=int, required=False, default = 50)
parser.add_argument('--repeats', type=int, required=False, default = 3)
parser.add_argument('--save_every', type=int, required=False, default = None)
parser.add_argument('--save_best_num', type=int, required=False, default = 1)
parser.add_argument('--eval_train', type=int, required=False, default = True)
parser.add_argument('--optimizer', type=str, required=False, default = "sgd")

args = parser.parse_args()
model_name = args.model_name
model_type = args.model_type
budget = args.budget
repeats = args.repeats
save_every = args.save_every 
eval_train = args.eval_train 
optimizer = args.optimizer.lower()
save_best_num = args.save_best_num 

os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"   
os.environ["CUDA_VISIBLE_DEVICES"]=args.gpu


def main():
    config_str = model_type + "." + model_name
    
    with open('hyperparameters.yaml') as f:
        hyperparameters = yaml.safe_load(f)    
    param_grid = hyperparameters[model_name]
    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
    savename = config(config_str + ".checkpoint")

    if not os.path.exists(savename):
        os.makedirs(savename)
    if not os.path.exists('{}/log/'.format(savename)):
        os.makedirs('{}/log/'.format(savename))

    exp = Experiment(optimizer,device, config_str, model_type, model_name, param_grid, save_best_num,savename, budget, repeats, eval_train = eval_train, save_every = save_every)

    df_search = exp.run()
    df_search.to_csv('{}/log/df_search.csv'.format(savename), index=False)
if __name__ == '__main__':
    main()