def train_new_data_with_model(): model = CoordRegression(n_locations=8) optimizer = optim.RMSprop(model.parameters(), lr=2.5e-4, alpha=0.9) model = torch.nn.DataParallel(model).cuda() from data_process_landmarks_hw import dataloader # 训练集 dataloader = dataloader for epoch in range(10): epoch_start = time.time() print("Epoch: {}/{}".format(epoch + 1, 10)) train_loss = 0.0 train_loss_cord = [] # forward pass model.train() # 训练 for i_batch, data in enumerate(dataloader): img, landmarks = data img = torch.tensor(img, dtype=torch.float32) img = img.to(device) landmarks = torch.tensor(landmarks / 64.0, dtype=torch.float32) landmarks = landmarks.to(device) # print("Ground-truth:", gt_hmap.shape) optimizer.zero_grad() # forward pass coords, heatmaps = model(img) # per-location euclidean losses euc_losses = dsntnn.euclidean_losses(coords, landmarks) # print("predict coords", coords, landmarks) # per-location regulation losses reg_losses = dsntnn.js_reg_losses(heatmaps, landmarks, sigma_t=1.0) # combine losses into an overall loss loss = dsntnn.average_loss(euc_losses + reg_losses) # calculate gradients optimizer.zero_grad() loss.backward() # update model parameters with RMSprop optimizer.step() train_loss_cord.append(loss) if i_batch % 20 == 19: print(loss, euc_losses, reg_losses) # break # print(loss) torch.save( model, 'models/' + 'landmarks' + '_model_new_data_8' + str(epoch) + '.pt') print(train_loss_cord)
def forward_2d_losses(self, out_var, target_var): target_xy = target_var.narrow(-1, 0, 2) losses = 0 for xy_hm, zy_hm, xz_hm in zip(self.xy_heatmaps, self.zy_heatmaps, self.xz_heatmaps): # Pixelwise heatmap loss. losses += self._calculate_pixelwise_loss(xy_hm, target_xy) # Calculated coordinate loss. actual_xy = self.heatmaps_to_coords(xy_hm, zy_hm, xz_hm)[..., :2] losses += euclidean_losses(actual_xy, target_xy) return losses
def forward_3d_losses(self, out_var, target_var): target_xyz = target_var.narrow(-1, 0, 3) losses = 0 target_xy = target_xyz.narrow(-1, 0, 2) target_zy = torch.cat([target_xyz.narrow(-1, 2, 1), target_xyz.narrow(-1, 1, 1)], -1) target_xz = torch.cat([target_xyz.narrow(-1, 0, 1), target_xyz.narrow(-1, 2, 1)], -1) for xy_hm, zy_hm, xz_hm in zip(self.xy_heatmaps, self.zy_heatmaps, self.xz_heatmaps): # Pixelwise heatmap loss. losses += self._calculate_pixelwise_loss(xy_hm, target_xy) losses += self._calculate_pixelwise_loss(zy_hm, target_zy) losses += self._calculate_pixelwise_loss(xz_hm, target_xz) # Calculated coordinate loss. actual_xyz = self.heatmaps_to_coords(xy_hm, zy_hm, xz_hm) losses += euclidean_losses(actual_xyz, target_xyz) return losses
def test_euclidean_mask(): output = torch.tensor([ [[0.0, 0.0], [1.0, 1.0], [0.0, 0.0]], [[1.0, 1.0], [0.0, 0.0], [0.0, 0.0]], ]) target = torch.tensor([ [[0.0, 0.0], [0.0, 0.0], [0.0, 0.0]], [[0.0, 0.0], [0.0, 0.0], [0.0, 0.0]], ]) mask = torch.tensor([ [1.0, 0.0, 1.0], [0.0, 1.0, 1.0], ]) expected = 0.0 actual = average_loss(euclidean_losses(output, target), mask) assert expected == actual.item()
def calc_coord_loss(coords, heatmaps, target_var, masks): # Per-location euclidean losses euc_losses = dsntnn.euclidean_losses( coords, target_var) # shape:[B, D, L, 2] batch, depth, locations, feature # Per-location regularization losses reg_losses = [] for i in range(heatmaps.shape[1]): hms = heatmaps[:, i] target = target_var[:, i] reg_loss = dsntnn.js_reg_losses(hms, target, sigma_t=1.0) reg_losses.append(reg_loss) reg_losses = torch.stack(reg_losses, 1) # reg_losses = dsntnn.js_reg_losses(heatmaps, target_var, sigma_t=1.0) # shape: [B, D, L, 7, 7] # Combine losses into an overall loss coord_loss = dsntnn.average_loss((euc_losses + reg_losses).squeeze(), mask=masks) return coord_loss
def forward(self, heatmaps, coords, targets, device): out = torch.Tensor([31, 31]).to(device) batch_size = coords.shape[0] n_stages = coords.shape[1] if len(targets.shape) != len(coords.shape): targets = torch.unsqueeze(targets, 1) targets = targets.repeat(1, n_stages, 1, 1) targets = (targets.div(255) * 2 + 1) / out - 1 losses = [] for i in range(batch_size): euc_loss = dsntnn.euclidean_losses(coords[i, :, :, :], targets[i, :, :, :]) reg_loss = dsntnn.js_reg_losses(heatmaps[i, :, :, :, :], targets[i, :, :, :], sigma_t=1.0) losses.append(dsntnn.average_loss(euc_loss + reg_loss)) return sum(losses) / batch_size
def test_mask(self): output = torch.Tensor([ [[0, 0], [1, 1], [0, 0]], [[1, 1], [0, 0], [0, 0]], ]) target = torch.Tensor([ [[0, 0], [0, 0], [0, 0]], [[0, 0], [0, 0], [0, 0]], ]) mask = torch.Tensor([ [1, 0, 1], [0, 1, 1], ]) expected = torch.Tensor([0]) actual = average_loss( euclidean_losses(Variable(output), Variable(target)), Variable(mask)) self.assertEqual(expected, actual.data)
def test_forward_and_backward(self): input_tensor = torch.Tensor([ [[3, 4], [3, 4]], [[3, 4], [3, 4]], ]) target = torch.Tensor([ [[0, 0], [0, 0]], [[0, 0], [0, 0]], ]) in_var = Variable(input_tensor, requires_grad=True) expected_loss = torch.Tensor([5]) actual_loss = average_loss(euclidean_losses(in_var, Variable(target))) expected_grad = torch.Tensor([ [[0.15, 0.20], [0.15, 0.20]], [[0.15, 0.20], [0.15, 0.20]], ]) actual_loss.backward() actual_grad = in_var.grad.data self.assertEqual(expected_loss, actual_loss.data) self.assertEqual(expected_grad, actual_grad)