(rank, world_size, conf)
| 42 | |
| 43 | |
| 44 | def train(rank, world_size, conf): |
| 45 | # create the log_file |
| 46 | os.makedirs(conf.exp.log_dir, exist_ok=True) |
| 47 | |
| 48 | # create the vis file |
| 49 | os.makedirs(conf.exp.vis_dir, exist_ok=True) |
| 50 | |
| 51 | setup(rank, world_size, conf) |
| 52 | |
| 53 | if dist.get_rank() == 0: |
| 54 | wandb.init(project='shape-matching', notes='weak baseline', config=conf) |
| 55 | wandb.define_metric("test/epoch/*", step_metric="test/epoch/epoch") |
| 56 | # wandb.define_metric("val/step/*", step_metric="val/step/step") |
| 57 | wandb.define_metric("train/network_B/*", step_metric="train/network_B/step") |
| 58 | wandb.define_metric("train/network_A/*", step_metric="train/network_A/step") |
| 59 | |
| 60 | data_features = ['src_pc', 'src_rot', 'src_trans', 'tgt_pc', 'tgt_rot', 'tgt_trans', 'partA_symmetry_type', 'partB_symmetry_type','predicted_partB_rotation', 'predicted_partB_position', 'predicted_partA_rotation', 'predicted_partA_position'] |
| 61 | |
| 62 | network_B = ShapeAssemblyNet_B_vnn(cfg=conf, data_features=data_features) |
| 63 | network_B.cuda(rank) |
| 64 | |
| 65 | # Load pretrained weights for network_B |
| 66 | network_B.load_state_dict(torch.load(os.path.join(conf.exp.log_dir, 'ckpts', 'network_B.pth'))) |
| 67 | |
| 68 | for param in network_B.parameters(): |
| 69 | param.requires_grad = False |
| 70 | |
| 71 | |
| 72 | network_A = ShapeAssemblyNet_A_vnn(cfg=conf, data_features=data_features) |
| 73 | network_A.cuda(rank) |
| 74 | |
| 75 | # Load pretrained weights for network_A |
| 76 | network_A.load_state_dict(torch.load(os.path.join(conf.exp.log_dir, 'ckpts', 'network_A.pth'))) |
| 77 | network_A = DDP(network_A, device_ids=[rank]) |
| 78 | |
| 79 | for param in network_A.parameters(): |
| 80 | param.requires_grad= False |
| 81 | |
| 82 | # Initialize train dataloader |
| 83 | train_data = OurDataset( |
| 84 | data_root_dir=conf.data.root_dir, |
| 85 | data_csv_file=conf.data.train_csv_file, |
| 86 | data_features=data_features, |
| 87 | num_points=conf.data.num_pc_points |
| 88 | ) |
| 89 | train_data.load_data() |
| 90 | |
| 91 | print('Len of Train Data: ', len(train_data)) |
| 92 | |
| 93 | # Initialize val dataloader |
| 94 | val_data = OurDataset( |
| 95 | data_root_dir=conf.data.root_dir, |
| 96 | data_csv_file=conf.data.val_csv_file, |
| 97 | data_features=data_features, |
| 98 | num_points=conf.data.num_pc_points |
| 99 | ) |
| 100 | val_data.load_data() |
| 101 |
nothing calls this directly
no test coverage detected