def main(config): # 选择使用模式 if config.mode == 'train': # 获取DataLoader train_loader = get_loader(config) # 计数, 找到可以创建文件夹的名字 run = 0 while os.path.exists("%s/run-%d" % (config.save_folder, run)): run += 1 # 创建保存模型文件夹 os.mkdir("%s/run-%d" % (config.save_folder, run)) os.mkdir("%s/run-%d/models" % (config.save_folder, run)) # 在config中保存文件夹名称 config.save_folder = "%s/run-%d" % (config.save_folder, run) # 开始进入training状态 train = Solver(train_loader, None, config) train.train() elif config.mode == 'test': # 选择数据集路径 config.test_root, config.test_list = get_test_info(config.sal_mode) # 获取DataLoader test_loader = get_loader(config, mode='test') # 如果文件路径不存在 if not os.path.exists(config.test_fold): os.mkdir(config.test_fold) # 开始进入testing状态 test = Solver(None, test_loader, config) test.test() else: # 非法输入 raise IOError("illegal input!!!")
def main(config): cudnn.benchmark = True if config.model_type not in ['U_Net','FTC_Net']: print('ERROR!! model_type should be selected in U_Net/FTC_Net') print('Your input for model_type was %s'%config.model_type) return if not os.path.exists(config.model_path): os.mkdir(config.model_path) if not os.path.exists(config.result_path): os.mkdir(config.result_path) config.result_path = os.path.join(config.result_path,config.model_type) if not os.path.exists(config.result_path): os.mkdir(config.result_path) train_loader = get_loader(config.save_path,'train',config.batch_size,config.num_workers) test_loader = get_loader(config.save_path, 'test', 1, config.num_workers) print('train dataset len %d' % (len(train_loader) * config.batch_size)) print('test dataset len %d' % len(test_loader)) results_path = config.result_path layer_results_path = os.path.join(results_path, 'layers') dst_results_path = os.path.join(results_path, 'dst_layer') att_result_path = os.path.join(results_path, 'attention_map') confi_result_path = os.path.join(results_path, 'confidence_map') if not os.path.exists(results_path): os.mkdir(results_path) if not os.path.exists(layer_results_path): os.mkdir(layer_results_path) os.mkdir(dst_results_path) if not os.path.exists(att_result_path): os.mkdir(att_result_path) if not os.path.exists(confi_result_path): os.mkdir(confi_result_path) if config.flag == 'FTC_Net': solver = Solver(config, train_loader, test_loader,) if config.mode == 'train': solver.train() elif config.mode == 'test': # os.mkdir(contour_results_path) # os.mkdir(attmap_path) # solver.test(layer_results_path,dst_results_path,att_result_path,config.is_save_attsmat,'./meter') solver.test(layer_results_path) elif config.mode == 'confi': solver.cal_confi(confi_result_path) else: solver.get_atts(layer_results_path,dst_results_path,att_result_path,config.is_save_attsmat)
def export(net, setup, system_env, system, exported_path="exported/output/net.pt", batch_size=1): print(colored("Trying to export to {}".format(exported_path), "red")) env_vox = torch.from_numpy( np.load('{}/data/{}_env_vox.npy'.format(system_env, system))).float() train_loader, test_loader = get_loader(system_env, system, batch_size=1, setup=setup) env_input = env_vox[train_loader.dataset[0:batch_size][0] [:, 0].long()].cuda() state_input = train_loader.dataset[0:batch_size][0][:, 1:].cuda() eval_net = copy.deepcopy(net) eval_net.eval() output = eval_net(state_input, env_input) print(colored("eval_output:", "blue")) print(output) eval_serilized_module = torch.jit.script(eval_net) print(colored("eval_serilized_output:", "blue")) print(eval_serilized_module(state_input, env_input)) print(colored("Tracing and Saving...", "blue")) traced_script_module = torch.jit.trace(net, (state_input, env_input)) serilized_module = torch.jit.script(net) serilized_output = serilized_module(state_input, env_input) traced_script_module.save(exported_path) print(colored("Exported to {}".format(exported_path), "red"))
def main(ae_output_size, state_size, lr, epochs, batch, system_env, system, setup, loss_type): mpnet = VAEMPNet(ae_input_size=32, ae_output_size=1024, in_channels=1, state_size=4, code_size=32) data_loaders = get_loader(system_env, system, batch_size=batch, setup=setup) train_network(network=mpnet, data_loaders=data_loaders, network_name="mpnet_vae", lr=lr, epochs=epochs, batch=batch, system_env=system_env, system=system, setup=setup, using_step_lr=True, step_size=50, gamma=0.9, loss_type=loss_type, weight_save_epochs=50)
def train(lr, epochs, batch, system, model, setup): mpnet = MPNet(ae_input_size=32, ae_output_size=128, in_channels=1, state_size=4, control_size=2) env_vox = torch.from_numpy( np.load('{}/{}_env_vox.npy'.format(system, model))).float() if torch.cuda.is_available(): mpnet = mpnet.cuda() env_vox = env_vox.cuda() #optimizer = torch.optim.Adagrad(mpnet.parameters(), lr=lr) #optimizer = torch.optim.SGD(mpnet.parameters(), lr=lr, momentum=0.9) optimizer = torch.optim.Adam(mpnet.parameters(), lr=lr) #scheduler = StepLR(optimizer, step_size=10, gamma=0.9) train_loader, test_loader = get_loader(system, model, batch_size=batch, setup=setup) for i in tqdm(range(epochs + 1)): print("epoch {}".format(i)) train_loss = [] mpnet.train() for data, label in tqdm(train_loader): ## prepare data if torch.cuda.is_available(): data = data.cuda() label = label.cuda() inputs = data[:, 1:] envs = env_vox[(data[:, 0]).long()] ## execute optimizer.zero_grad() output = mpnet(inputs, envs) loss = F.mse_loss(output, label) loss.backward() optimizer.step() train_loss.append(loss.item()) # break #scheduler.step(i) print("train loss:{}".format(np.mean(train_loss))) for data, label in test_loader: mpnet.eval() eval_loss = [] with torch.no_grad(): if torch.cuda.is_available(): data = data.cuda() label = label.cuda() inputs = data[:, 1:] envs = env_vox[(data[:, 0]).long()] output = mpnet(inputs, envs) loss = F.mse_loss(output, label) eval_loss.append(loss.item()) print("eval loss:{}".format(np.mean(eval_loss))) if i % 5 == 0: Path("output/{}/{}".format(model, setup)).mkdir(parents=True, exist_ok=True) torch.save(mpnet.state_dict(), "output/{}/{}/ep{}".format(model, setup, i))
def main(config): if config.mode == 'train': train_loader = get_loader(config) run = 0 while os.path.exists("%s/run-%d" % (config.save_folder, run)): run += 1 os.mkdir("%s/run-%d" % (config.save_folder, run)) os.mkdir("%s/run-%d/models" % (config.save_folder, run)) config.save_folder = "%s/run-%d" % (config.save_folder, run) train = Solver(train_loader, None, config) train.train() elif config.mode == 'test': config.test_root, config.test_list = get_test_info(config.sal_mode) test_loader = get_loader(config, mode='test') if not os.path.exists(config.test_fold): os.mkdir(config.test_fold) test = Solver(None, test_loader, config) test.test() else: raise IOError("illegal input!!!")
def main(lr, num_epoch, batch_size, loss_type, network_tag): net = CNNWrapper(loss_type=loss_type, save_path="./output", device="cuda", tag=network_tag) data_loaders = get_loader(batch_size=batch_size) tvs = TrainValScheduler(network=net, data_loaders=data_loaders, lr=1e-3, num_epoch=1, batch_size=200, using_step_lr=False, step_size=50, gamma=0.9, weight_save_epochs=50, logging_path="log/", optimizer_type="Adadelta") tvs.train_val()
def main(ae_output_size, state_size, lr, epochs, batch, system_env, system, setup, loss_type, lr_step_size, aug, network_name): # mpnet_module = importlib.import_module('.mpnet_{}'.format(system), package=".networks") try: if system == 'cartpole_obs': if network_name == 'mpnet': from networks.mpnet_cartpole_obs import MPNet elif network_name == 'mpnet_branch': from networks.mpnet_cartpole_obs_branch import MPNet state_dim = 4 in_channels = 1 elif system == 'acrobot_obs': from networks.mpnet_acrobot_obs import MPNet state_dim = 4 in_channels = 1 elif system == 'quadrotor_obs': from networks.mpnet_quadrotor_obs import MPNet state_dim = 13 in_channels = 32 elif system == 'car_obs': from networks.mpnet_car_obs import MPNet state_dim = 3 in_channels = 1 except ModuleNotFoundError: print("Unrecognized model name") raise mpnet = MPNet(ae_input_size=32, ae_output_size=ae_output_size, in_channels=in_channels, state_size=state_dim) data_loaders = get_loader(system_env, system, batch_size=batch, setup=setup) train_network(network=mpnet, data_loaders=data_loaders, network_name=network_name, lr=lr, epochs=epochs, batch=batch, system_env=system_env, system=system, setup=setup, using_step_lr=True, step_size=lr_step_size, gamma=0.9, loss_type=loss_type, weight_save_epochs=25, aug=aug)
def main(config): test_loader = get_loader(test_mode=config.test_mode, sal_mode=config.sal_mode) if not os.path.exists(config.test_fold): os.mkdir(config.test_fold) test = Solver(test_loader, config) test.test(test_mode=config.test_mode)
def train(lr, epochs, batch, system, model, setup): mpnet = VAEMPNet(ae_input_size=32, ae_output_size=1024, in_channels=1, state_size=4, code_size=32) env_vox = torch.from_numpy( np.load('{}/{}_env_vox.npy'.format(system, model))).float() if torch.cuda.is_available(): mpnet = mpnet.cuda() env_vox = env_vox.cuda() optimizer = torch.optim.Adam(mpnet.parameters(), lr=lr) scheduler = StepLR(optimizer, step_size=50, gamma=0.9) train_loader, test_loader = get_loader(system, model, batch_size=batch, setup=setup) with tqdm(range(epochs + 1), total=epochs + 1) as pbar: for i in range(epochs + 1): #print("epoch {}".format(i)) train_loss = [] mpnet.train() for data, label in tqdm(train_loader): ## prepare data if torch.cuda.is_available(): data = data.cuda() label = label.cuda() inputs = data[:, 1:] envs = env_vox[(data[:, 0]).long()] ## execute optimizer.zero_grad() output = mpnet(inputs, envs) # loss = F.mse_loss(output, label) loss = F.l1_loss(output, label) loss.backward() # print(loss.item()) # train_loss.append(loss.item()) optimizer.step() # break scheduler.step(i) # print("train loss:{}".format(np.mean(train_loss))) mpnet.eval() for data, label in test_loader: eval_loss = [] with torch.no_grad(): if torch.cuda.is_available(): data = data.cuda() label = label.cuda() inputs = data[:, 1:] envs = env_vox[(data[:, 0]).long()] output = mpnet(inputs, envs) # loss = F.mse_loss(output, label) loss = F.l1_loss(output, label) eval_loss.append(loss.item()) #print("eval loss:{}".format(np.mean(eval_loss))) pbar.set_postfix( {'eval_loss': '{0:1.5f}'.format(np.mean(eval_loss))}) if i % 50 == 0: Path("output/{}/{}".format(model, setup)).mkdir(parents=True, exist_ok=True) torch.save(mpnet.state_dict(), "output/{}/{}/ep{}".format(model, setup, i)) pbar.update(1)
parser.add_argument('--train_root', type=str, default='') parser.add_argument('--train_list', type=str, default='') # Testing settings parser.add_argument('--model', type=str, default='./dataset/pretrained/final.pth') # Snapshot parser.add_argument('--test_fold', type=str, default='./results') # Test results saving folder parser.add_argument('--sal_mode', type=str, default='g') # Test image dataset # Misc parser.add_argument('--mode', type=str, default='test', choices=['train', 'test']) config = parser.parse_args() if not os.path.exists(config.save_folder): os.mkdir(config.save_folder) # Get test set info config.test_root = './data/game/' config.test_list = './data/game/test.lst' config.cuda = False test_loader = get_loader(config, mode='test') if not os.path.exists(config.test_fold): os.mkdir(config.test_fold) test = Solver(None, test_loader, config) #test.test() ##################################################### # (2) load pretrained weights ##################################################### #------------------------------------------ import io import numpy as np from torch import nn import torch.utils.model_zoo as model_zoo import torch.onnx import torch.nn as nn