core training loop for featurenet
(args, targets, rgbs, poses, feat_model, dset_size, FeatureLoss, optimizer, hwf)
| 100 | args = parser.parse_args() |
| 101 | |
| 102 | def train_on_batch(args, targets, rgbs, poses, feat_model, dset_size, FeatureLoss, optimizer, hwf): |
| 103 | ''' core training loop for featurenet''' |
| 104 | feat_model.train() |
| 105 | H, W, focal = hwf |
| 106 | H, W = int(H), int(W) |
| 107 | if args.freezeBN: |
| 108 | feat_model = freeze_bn_layer_train(feat_model) |
| 109 | |
| 110 | train_loss_epoch = [] |
| 111 | select_inds = np.random.choice(dset_size, size=[dset_size], replace=False) # (N_rand,) |
| 112 | |
| 113 | batch_size=args.featurenet_batch_size # manual setting, use smaller batch size like featurenet_batch_size = 4 if OOM |
| 114 | if dset_size % batch_size == 0: |
| 115 | N_iters = dset_size//batch_size |
| 116 | else: |
| 117 | N_iters = dset_size//batch_size + 1 |
| 118 | i_batch = 0 |
| 119 | |
| 120 | for i in range(0, N_iters): |
| 121 | if i_batch + batch_size > dset_size: |
| 122 | i_batch = 0 |
| 123 | break |
| 124 | i_inds = select_inds[i_batch:i_batch+batch_size] |
| 125 | i_batch = i_batch + batch_size |
| 126 | |
| 127 | # convert input shape to [B, 3, H, W] |
| 128 | target_in = targets[i_inds].clone().permute(0,3,1,2).to(device) |
| 129 | rgb_in = rgbs[i_inds].clone().permute(0,3,1,2).to(device) |
| 130 | pose = poses[i_inds].clone().reshape(batch_size, 12).to(device) |
| 131 | pose = torch.cat([pose, pose]) # double gt pose tensor |
| 132 | |
| 133 | features, predict_pose = feat_model(torch.cat([target_in, rgb_in]), True, upsampleH=H, upsampleW=W) # features: (1, [2, B, C, H, W]) |
| 134 | |
| 135 | # get features_target and features_rgb |
| 136 | if args.DFNet: |
| 137 | features_target = features[0] # [3, B, C, H, W] |
| 138 | features_rgb = features[1] |
| 139 | else: |
| 140 | features_target = features[0][0] |
| 141 | features_rgb = features[0][1] |
| 142 | |
| 143 | # svd, seems not very benificial here, therefore removed |
| 144 | |
| 145 | if args.poselossonly: |
| 146 | loss_pose = PoseLoss(args, predict_pose, pose, device) # target |
| 147 | loss = loss_pose |
| 148 | elif args.featurelossonly: # Not good. To be removed later |
| 149 | loss_f = FeatureLoss(features_rgb, features_target) |
| 150 | loss = loss_f |
| 151 | else: |
| 152 | loss_pose = PoseLoss(args, predict_pose, pose, device) # target |
| 153 | if args.tripletloss: |
| 154 | loss_f = triplet_loss_hard_negative_mining_plus(features_rgb, features_target, margin=args.triplet_margin) |
| 155 | else: |
| 156 | loss_f = FeatureLoss(features_rgb, features_target) |
| 157 | loss = loss_pose + loss_f |
| 158 | |
| 159 | loss.backward() |
no test coverage detected