def main(args): os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" # see issue #152 os.environ["CUDA_VISIBLE_DEVICES"] = "0" importer = fileImport() env_data_path = args.env_data_path path_data_path = args.path_data_path pcd_data_path = args.pointcloud_data_path envs = importer.environments_import(env_data_path + args.envs_file) print("Loading obstacle data...\n") dataset_train, targets_train, pc_inds_train, obstacles = load_dataset_end2end( envs, path_data_path, pcd_data_path, args.path_data_file, importer, NP=1000) print("Loaded dataset, targets, and pontcloud obstacle vectors: ") print(str(len(dataset_train)) + " " + str(len(targets_train)) + " " + str(len(pc_inds_train))) print("\n") if not os.path.exists(args.trained_model_path): os.makedirs(args.trained_model_path) # Build the models mlp = MLP(args.mlp_input_size, args.mlp_output_size) encoder = Encoder(args.enc_input_size, args.enc_output_size) if torch.cuda.is_available(): encoder.cuda() mlp.cuda() # Loss and Optimizer criterion = nn.MSELoss() params = list(encoder.parameters())+list(mlp.parameters()) optimizer = torch.optim.Adagrad(params, lr=args.learning_rate) total_loss = [] epoch = 1 sm = 90 # start saving models after 100 epochs print("Starting epochs...\n") # epoch=1 done = False for epoch in range(args.num_epochs): # while (not done) start = time.time() print("epoch" + str(epoch)) avg_loss = 0 for i in range(0, len(dataset_train), args.batch_size): # Forward, Backward and Optimize # zero gradients encoder.zero_grad() mlp.zero_grad() # convert to pytorch tensors and Varialbes bi, bt, bobs = get_input( i, dataset_train, targets_train, pc_inds_train, obstacles, args.batch_size) bi = to_var(bi) bt = to_var(bt) bobs = to_var(bobs) # forward pass through encoder h = encoder(bobs) # concatenate encoder output with dataset input inp = torch.cat((bi, h), dim=1) # forward pass through mlp bo = mlp(inp) # compute overall loss and backprop all the way loss = criterion(bo, bt) avg_loss = avg_loss+loss.data loss.backward() optimizer.step() print("--average loss:") print(avg_loss/(len(dataset_train)/args.batch_size)) total_loss.append(avg_loss/(len(dataset_train)/args.batch_size)) # Save the models if epoch == sm: print("\nSaving model\n") print("time: " + str(time.time() - start)) torch.save(encoder.state_dict(), os.path.join( args.trained_model_path, 'cae_encoder_'+str(epoch)+'.pkl')) torch.save(total_loss, 'total_loss_'+str(epoch)+'.dat') model_path = 'mlp_PReLU_ae_dd'+str(epoch)+'.pkl' torch.save(mlp.state_dict(), os.path.join( args.trained_model_path, model_path)) if (epoch != 1): sm = sm+10 # save model after every 50 epochs from 100 epoch ownwards torch.save(total_loss, 'total_loss.dat') model_path = 'mlp_PReLU_ae_dd_final.pkl' torch.save(mlp.state_dict(), os.path.join(args.trained_model_path, model_path))
num_workers=8) """ Model settings """ mesh = kal.rep.TriangleMesh.from_obj('386.obj') if args.device == "cuda": mesh.cuda() initial_verts = mesh.vertices.clone() model = Encoder(4, 5, args.batchsize, 137, mesh.vertices.shape[0]).to(args.device) loss_fn = kal.metrics.point.chamfer_distance loss_edge = kal.metrics.mesh.edge_length optimizer = optim.Adam(model.parameters(), lr=args.lr) # Create log directory, if it doesn't already exist args.logdir = os.path.join(args.logdir, args.expid) if not os.path.isdir(args.logdir): os.makedirs(args.logdir) print('Created dir:', args.logdir) # Log all commandline args with open(os.path.join(args.logdir, 'args.txt'), 'w') as f: json.dump(args.__dict__, f, indent=2) class Engine(object): """Engine that runs training and inference. Args