def save_network(model, args, folder, delimiter=""):
file_utils.create_folder(folder)
if args.network_model.lower() == "stein":
file_name = "./" + folder + "/svgd_particles_" + str(
args.num_particles) + "_dropout_rate_" + str(
args.dropout_rate) + "_reg_" + str(
args.regularization) + "_lat_" + str(
args.latent_dim) + delimiter + ".model"
tc.save(model.state_dict(), file_name)
elif args.network_model.lower() == "mc_dropout":
file_name = "./" + folder + "/mc_dropout_dropout_rate_" + str(
args.dropout_rate) + "_lat_" + str(
args.latent_dim) + delimiter + ".model"
tc.save(model.state_dict(), file_name)
def save_test_output(
predictions,
object_name,
observed_pc,
object_pose,
test_case,
storage_folder,
network_model,
save_voxel_grid=True,
save_samples=False,
save_mesh=False,
):
if network_model.lower() != "varley":
mean_voxel = shape_completion_utils.calculate_mean_voxel_of_samples(
predictions)
else:
mean_voxel = predictions
if save_voxel_grid:
voxel_folder = file_utils.create_folder(storage_folder + "/voxels/" +
test_case + "/" + object_name +
"/")
shape_completion_utils.save_voxel_grid(mean_voxel,
object_name + "_mean_shape",
voxel_folder)
binvox_folder = file_utils.create_folder(storage_folder + "/binvox/" +
test_case + "/" +
object_name + "/")
save_binvox_input_output(observed_pc, mean_voxel, binvox_folder)
if save_mesh:
mesh_folder = file_utils.create_folder(storage_folder + "/meshes/" +
test_case + "/" + object_name +
"/")
shape_completion_utils.cnn_and_pc_to_mesh(
observed_pc, mean_voxel, mesh_folder,
object_name + "_mean_shape.ply", object_pose)
if save_samples and network_model.lower() != "varley":
predictions = predictions
for sample in range(predictions.shape[0]):
sample_name = object_name + "_sample_" + str(sample)
if save_voxel_grid:
shape_completion_utils.save_voxel_grid(predictions[sample],
sample_name,
voxel_folder)
if save_mesh:
shape_completion_utils.cnn_and_pc_to_mesh(
observed_pc, predictions[sample], mesh_folder,
sample_name + ".ply", object_pose)
def train_model(args):
network_utils.setup_env(args.use_cuda)
model = network_utils.setup_network(args)
model.set_mode(train=True)
net_dump_folder = "networks/"
file_utils.create_folder("%scheckpoints/" % (net_dump_folder))
start_epoch = 0
if args.net_recover:
start_epoch = network_utils.recover_network(args, model)
else:
network_utils.save_network(model, args,
net_dump_folder + "checkpoints/",
"_epoch_" + str(0))
dataloader = network_utils.load_dataset(args.num_workers,
args.batch_size,
args.mode,
debug=args.debug)
print("Starting training (%d samples, %d epochs)" %
(len(dataloader.dataset), args.num_epochs))
file_utils.create_folder(net_dump_folder + "training_results/")
for epoch in range(start_epoch + 1, args.num_epochs + 1):
running_loss = []
print("EPOCH NUMBER " + str(epoch))
for idx, data in enumerate(dataloader):
inputs, targets = network_utils.get_train_data(data, args.use_cuda)
num_data = inputs.size(0)
loss = model.train_network(inputs, targets)
running_loss.append(loss.item() / num_data)
if (epoch + 1) % args.checkpoint_interval == 0:
network_utils.save_network(model, args,
net_dump_folder + "checkpoints/",
"_epoch_" + str(int(epoch)))
file_utils.log_data_to_file(
np.mean(running_loss), net_dump_folder + "training_results/loss_" +
args.network_model + ".txt")
model.eval()
print("Finished training")
# Saving final network to file
print("Saving final network parameters")
network_utils.save_network(model, args, net_dump_folder + "final_network/")
tc.cuda.empty_cache()