# ---------------------- Save checkpoint --------------------- #
if epoch % 5 == 0 and not save_ckpt == None:
torch.save(
{
'epoch': epoch,
'D_state_dict': self.D.state_dict(),
'G_state_dict': self.G.state_dict(),
'D_loss': loss_log['D_loss'],
'G_loss': loss_log['G_loss'],
'FPD': metric['FPD']
}, save_ckpt + str(epoch) + '_' + class_name + '.pt')
# print('Checkpoint is saved.')
if __name__ == '__main__':
args = Arguments().parser().parse_args()
args.device = torch.device(
'cuda:' + str(args.gpu) if torch.cuda.is_available() else 'cpu')
torch.cuda.set_device(args.device)
SAVE_CHECKPOINT = args.ckpt_path + args.ckpt_save if args.ckpt_save is not None else None
LOAD_CHECKPOINT = args.ckpt_load if args.ckpt_load is not None else None
RESULT_PATH = args.result_path + args.result_save
model = TreeGAN(args)
model.run(save_ckpt=SAVE_CHECKPOINT,
load_ckpt=LOAD_CHECKPOINT,
result_path=RESULT_PATH)
gt_data = torch.Tensor([])
for _iter, data in enumerate(dataLoader):
point, partial, index = data
gt_data = torch.cat((gt_data, point), 0)
ref_pcs = gt_data.detach().cpu().numpy()
sample_pcs = fake_pcs.detach().cpu().numpy()
tic = time.time()
mmd, matched_dists, dist_mat = MMD_batch(sample_pcs,
ref_pcs,
batch_size,
normalize=normalize,
use_EMD=use_EMD,
device=args.device)
toc = time.time()
if verbose:
print('-----MMD-EMD: {:5.3f}'.format(mmd * 100))
if __name__ == '__main__':
args = Arguments(stage='eval_treegan').parser().parse_args()
args.device = torch.device('cuda')
assert args.eval_treegan_mode in [
"MMD", "FPD", "save", "generate_fpd_stats"
]
if args.eval_treegan_mode == "generate_fpd_stats":
script_create_fpd_stats(args)
else:
test(args, mode=args.eval_treegan_mode)
