TEST_DIF_DATA_PATH, 'rb', ) as f: test_dif_data = pickle.load(f) CUDA_VISIBLE_DEVICES = '1' print('Current GPU index: ' + CUDA_VISIBLE_DEVICES) os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" os.environ["CUDA_VISIBLE_DEVICES"] = CUDA_VISIBLE_DEVICES if torch.cuda.is_available(): cuda_flag = True else: cuda_flag = False maze_core = MazeCore() npi = NPI(maze_core, CONFIG) if cuda_flag: npi = npi.cuda() print_net(npi) parameters = filter(lambda p: p.requires_grad, npi.parameters()) optimizer = torch.optim.SGD(parameters, lr=0.0001, momentum=0.9, nesterov=True, weight_decay=5e-4) # optimizer = torch.optim.Adam(parameters, lr=0.0001, betas=(0.5, 0.999)) lr_schedulers = torch.optim.lr_scheduler.CosineAnnealingLR( optimizer, T_max=max_num_epochs)
with open( TEST_DIF_DATA_PATH, 'rb', ) as f: test_dif_data = pickle.load(f) if torch.cuda.is_available(): cuda_flag = True CUDA_VISIBLE_DEVICES = '0' print('Current GPU index: ' + CUDA_VISIBLE_DEVICES) os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" os.environ["CUDA_VISIBLE_DEVICES"] = CUDA_VISIBLE_DEVICES else: cuda_flag = False maze_core = MazeCore(hidden_dim=100) npi = NPI(maze_core, CONFIG) npi.load_networks(load_path) if cuda_flag: npi = npi.cuda() if cuda_flag: npi = npi.cuda() print_net(npi) mode = 'test_same' cur_same_results = run_epoch(npi, mode, test_same_data, writer) same_def_loss, same_total_loss, same_ter_accs, same_pro_accs = cur_same_results