def main(): print("Start... => main.py PID: %s" % (os.getpid())) opt = get_opt() # torch.manual_seed(opt.seed) # random.seed(opt.seed) # np.random.seed(opt.seed) # print(opt.gpus) # if opt.gpus and torch.cuda.is_available(): # print("opt.gpus: ", opt.gpus) # gpu_list = [int(k) for k in opt.gpus.split(",")] # print("gpu_list: ", gpu_list) # print("gpu_list[0]: ", gpu_list[0]) # print("type(gpu_list[0]): ", type(gpu_list[0])) # os.environ["CUDA_VISIBLE_DEVICES"] = opt.gpus # all_dict = load_dict() if opt.train_mode == "train": data_loader, all_dict = load_data(opt) print('Loaded dataset sucessfully.') model = create_code_retrieval_model(opt, all_dict) print('created model..') cos_ranking_loss = CosRankingLoss( opt).cuda() if torch.cuda.is_available() else CosRankingLoss(opt) optim = torch.optim.Adam(model.parameters(), opt.lr) print('created criterion and optim') # standard trainer = CodeRetrievalTrainer(model, data_loader, all_dict, opt) # begin to training !!! trainer.train(cos_ranking_loss, optim, opt.train_epoch) if opt.train_mode == "query": data_loader, all_dict, code_body, func_name = load_data( opt, model='n_query') model = create_code_retrieval_model(opt, all_dict) model.load_state_dict( torch.load('./model/Models/model_0.pt', map_location=lambda storage, loc: storage)) print('Loaded dataset sucessfully.') code_searcher = CodereRerievalSearcher(10, code_body, func_name, all_dict, opt, model, data_loader) result = code_searcher.search('set working directory') code_searcher.save_result()
import torchvision.datasets as dset import torchvision.transforms as transforms from torchvision.utils import save_image from torch.autograd import Variable from tensorboardX import SummaryWriter import models.dcgan_model as dcgan from utils import save_checkpoint, load_checkpoint, save_checkpoint_wo_step from utils import board_add_image, board_add_images from utils import save_image_historys_gif from utils import inception_score from opt import get_opt opt = get_opt() print(opt) # ======= # device # ======= cuda = True if torch.cuda.is_available() else False if (opt.device == "gpu"): if (cuda): device = torch.device("cuda") else: device = torch.device("cpu") else: device = torch.device("cpu") # ===================
def main(): print("Start... => main.py PID: %s" % (os.getpid())) opt = get_opt() torch.manual_seed(opt.seed) random.seed(opt.seed) np.random.seed(opt.seed) if opt.gpus: print("opt.gpus: ", opt.gpus) gpu_list = [int(k) for k in opt.gpus.split(",")] print("gpu_list: ", gpu_list) print("gpu_list[0]: ", gpu_list[0]) print("type(gpu_list[0]): ", type(gpu_list[0])) os.environ["CUDA_VISIBLE_DEVICES"] = opt.gpus all_dict = load_dict(opt) if opt.train_mode == "train": if opt.validation_with_metric: train_dataset, train_dataloader, val_dataset, val_dataloader, \ test_dataset, test_dataloader, query_dataset, query_dataloader = load_data(opt, all_dict) else: train_dataset, train_dataloader, val_dataset, val_dataloader = load_data( opt, all_dict) elif opt.train_mode == "test": test_dataset, test_dataloader, query_dataset, query_dataloader = load_data( opt, all_dict) print('Loaded dataset sucessfully.') if opt.train_mode == "train": model = create_model_code_retrieval(opt, train_dataset, all_dict) print('created model..') cos_ranking_loss = CosRankingLoss(opt).cuda() optim = torch.optim.Adam(model.parameters(), opt.lr) print('created criterion and optim') if opt.validation_with_metric: pretrainer = CodeRetrievalTrainer(model, train_dataset, train_dataloader, val_dataset, val_dataloader, all_dict, opt, metric_data=[ test_dataset, test_dataloader, query_dataset, query_dataloader ]) else: pretrainer = CodeRetrievalTrainer(model, train_dataset, train_dataloader, val_dataset, val_dataloader, all_dict, opt) pretrainer.train(cos_ranking_loss, optim, opt.pretrain_sl_epoch) elif opt.train_mode == "test": model = create_model_code_retrieval(opt, test_dataset, all_dict) print('created model..') evaluator = CodeRetrievalEvaluator(model=model, dataset_list=[ test_dataset, test_dataloader, query_dataset, query_dataloader ], flag_for_val=False, all_dict=all_dict, opt=opt) evaluator.retrieval(pred_file=opt.retrieval_pred_file) if opt.use_val_as_codebase: evaluator.eval_retrieval_json_result( pred_file=opt.retrieval_pred_file)