def test(args, net, test_loader, boardio, textio):
test_loss, test_cycle_loss, \
test_mse_ab, test_mae_ab, test_mse_ba, test_mae_ba, test_rotations_ab, test_translations_ab, \
test_rotations_ab_pred, \
test_translations_ab_pred, test_rotations_ba, test_translations_ba, test_rotations_ba_pred, \
test_translations_ba_pred, test_eulers_ab, test_eulers_ba = test_one_epoch(args, net, test_loader)
test_rmse_ab = np.sqrt(test_mse_ab)
test_rmse_ba = np.sqrt(test_mse_ba)
test_rotations_ab_pred_euler = npmat2euler(test_rotations_ab_pred)
test_r_mse_ab = np.mean(
(test_rotations_ab_pred_euler - np.degrees(test_eulers_ab))**2)
test_r_rmse_ab = np.sqrt(test_r_mse_ab)
test_r_mae_ab = np.mean(
np.abs(test_rotations_ab_pred_euler - np.degrees(test_eulers_ab)))
test_t_mse_ab = np.mean(
(test_translations_ab - test_translations_ab_pred)**2)
test_t_rmse_ab = np.sqrt(test_t_mse_ab)
test_t_mae_ab = np.mean(
np.abs(test_translations_ab - test_translations_ab_pred))
test_rotations_ba_pred_euler = npmat2euler(test_rotations_ba_pred, 'xyz')
test_r_mse_ba = np.mean(
(test_rotations_ba_pred_euler - np.degrees(test_eulers_ba))**2)
test_r_rmse_ba = np.sqrt(test_r_mse_ba)
test_r_mae_ba = np.mean(
np.abs(test_rotations_ba_pred_euler - np.degrees(test_eulers_ba)))
test_t_mse_ba = np.mean(
(test_translations_ba - test_translations_ba_pred)**2)
test_t_rmse_ba = np.sqrt(test_t_mse_ba)
test_t_mae_ba = np.mean(
np.abs(test_translations_ba - test_translations_ba_pred))
textio.cprint('==FINAL TEST==')
textio.cprint('A--------->B')
textio.cprint(
'EPOCH:: %d, Loss: %f, Cycle Loss: %f, MSE: %f, RMSE: %f, MAE: %f, rot_MSE: %f, rot_RMSE: %f, '
'rot_MAE: %f, trans_MSE: %f, trans_RMSE: %f, trans_MAE: %f' %
(-1, test_loss, test_cycle_loss, test_mse_ab, test_rmse_ab,
test_mae_ab, test_r_mse_ab, test_r_rmse_ab, test_r_mae_ab,
test_t_mse_ab, test_t_rmse_ab, test_t_mae_ab))
textio.cprint('B--------->A')
textio.cprint(
'EPOCH:: %d, Loss: %f, MSE: %f, RMSE: %f, MAE: %f, rot_MSE: %f, rot_RMSE: %f, '
'rot_MAE: %f, trans_MSE: %f, trans_RMSE: %f, trans_MAE: %f' %
(-1, test_loss, test_mse_ba, test_rmse_ba, test_mae_ba, test_r_mse_ba,
test_r_rmse_ba, test_r_mae_ba, test_t_mse_ba, test_t_rmse_ba,
test_t_mae_ba))
def _train_one_epoch(self, epoch, train_loader, opt, temp):
self.train()
total_loss = 0
rotations_ab = []
translations_ab = []
rotations_ab_pred = []
translations_ab_pred = []
eulers_ab = []
num_examples = 0
for data in tqdm(train_loader):
src, tgt, rotation_ab, translation_ab, rotation_ba, translation_ba, euler_ab, euler_ba = [d.cuda()
for d in data]
loss, rotation_ab_pred, translation_ab_pred = self._train_one_batch(src, tgt, rotation_ab, translation_ab,
opt, temp)
batch_size = src.size(0)
num_examples += batch_size
total_loss = total_loss + loss * batch_size
rotations_ab.append(rotation_ab.detach().cpu().numpy())
translations_ab.append(translation_ab.detach().cpu().numpy())
rotations_ab_pred.append(rotation_ab_pred.detach().cpu().numpy())
translations_ab_pred.append(translation_ab_pred.detach().cpu().numpy())
eulers_ab.append(euler_ab.cpu().numpy())
avg_loss = total_loss / num_examples
# (num_examples, 3, 3)
rotations_ab = np.concatenate(rotations_ab, axis=0)
translations_ab = np.concatenate(translations_ab, axis=0)
rotations_ab_pred = np.concatenate(rotations_ab_pred, axis=0)
translations_ab_pred = np.concatenate(translations_ab_pred, axis=0)
eulers_ab = np.degrees(np.concatenate(eulers_ab, axis=0))
eulers_ab_pred = npmat2euler(rotations_ab_pred)
r_ab_mse = np.mean((eulers_ab - eulers_ab_pred) ** 2)
r_ab_rmse = np.sqrt(r_ab_mse)
r_ab_mae = np.mean(np.abs(eulers_ab - eulers_ab_pred))
rot = np.matmul(rotations_ab_pred.transpose(0, 2, 1), rotations_ab)
rot_trace = rot[:, 0, 0] + rot[:, 1, 1] + rot[:, 2, 2]
residual_rotdeg = np.arccos(np.clip(0.5 * (rot_trace - 1), a_min=-1.0, a_max=1.0)) * 180.0 / np.pi
residual_rotdeg = np.mean(residual_rotdeg)
t_ab_mse = np.mean((translations_ab - translations_ab_pred) ** 2)
t_ab_rmse = np.sqrt(t_ab_mse)
t_ab_mae = np.mean(np.abs(translations_ab - translations_ab_pred))
r_ab_r2_score = r2_score(eulers_ab, eulers_ab_pred)
t_ab_r2_score = r2_score(translations_ab, translations_ab_pred)
info = {'arrow': 'A->B',
'epoch': epoch,
'stage': 'train',
'loss': avg_loss,
'r_ab_mse': r_ab_mse,
'r_ab_rmse': r_ab_rmse,
'r_ab_mae': r_ab_mae,
't_ab_mse': t_ab_mse,
't_ab_rmse': t_ab_rmse,
't_ab_mae': t_ab_mae,
'r_ab_r2_score': r_ab_r2_score,
't_ab_r2_score': t_ab_r2_score,
'residual_rotdeg': residual_rotdeg,
'temperature': temp}
self.logger.write(info)
return info
def _test_one_epoch(self, epoch, test_loader, temp):
self.eval()
total_loss = 0
rotations_ab = []
translations_ab = []
rotations_ab_pred = []
translations_ab_pred = []
eulers_ab = []
num_examples = 0
for data in tqdm(test_loader):
src, tgt, rotation_ab, translation_ab, rotation_ba, translation_ba, euler_ab, euler_ba = [
d.cuda() for d in data
]
loss, rotation_ab_pred, translation_ab_pred = self._test_one_batch(
src, tgt, rotation_ab, translation_ab, temp)
batch_size = src.size(0)
num_examples += batch_size
total_loss = total_loss + loss * batch_size
rotations_ab.append(rotation_ab.detach().cpu().numpy())
translations_ab.append(translation_ab.detach().cpu().numpy())
rotations_ab_pred.append(rotation_ab_pred.detach().cpu().numpy())
translations_ab_pred.append(
translation_ab_pred.detach().cpu().numpy())
eulers_ab.append(euler_ab.cpu().numpy())
avg_loss = total_loss / num_examples
rotations_ab = np.concatenate(rotations_ab, axis=0)
translations_ab = np.concatenate(translations_ab, axis=0)
rotations_ab_pred = np.concatenate(rotations_ab_pred, axis=0)
translations_ab_pred = np.concatenate(translations_ab_pred, axis=0)
eulers_ab = np.degrees(np.concatenate(eulers_ab, axis=0))
eulers_ab_pred = npmat2euler(rotations_ab_pred)
r_ab_mse = np.mean((eulers_ab - eulers_ab_pred)**2)
r_ab_rmse = np.sqrt(r_ab_mse)
r_ab_mae = np.mean(np.abs(eulers_ab - eulers_ab_pred))
t_ab_mse = np.mean((translations_ab - translations_ab_pred)**2)
t_ab_rmse = np.sqrt(t_ab_mse)
t_ab_mae = np.mean(np.abs(translations_ab - translations_ab_pred))
r_ab_r2_score = r2_score(eulers_ab, eulers_ab_pred)
t_ab_r2_score = r2_score(translations_ab, translations_ab_pred)
info = {
'arrow': 'A->B',
'epoch': epoch,
'stage': 'test',
'loss': avg_loss,
'r_ab_mse': r_ab_mse,
'r_ab_rmse': r_ab_rmse,
'r_ab_mae': r_ab_mae,
't_ab_mse': t_ab_mse,
't_ab_rmse': t_ab_rmse,
't_ab_mae': t_ab_mae,
'r_ab_r2_score': r_ab_r2_score,
't_ab_r2_score': t_ab_r2_score,
'temperature': temp
}
self.logger.write(info)
return info
def __getitem__(self, idx):
file0 = os.path.join(self.root, self.files[idx][0])
file1 = os.path.join(self.root, self.files[idx][1])
data0 = np.load(file0)
data1 = np.load(file1)
xyz0 = data0["pcd"]
xyz1 = data1["pcd"]
if self.random_rotation:
T0 = sample_random_trans(xyz0, self.randg, self.rotation_range)
T1 = sample_random_trans(xyz1, self.randg, self.rotation_range)
trans = T1 @ np.linalg.inv(T0)
xyz0 = self.apply_transform(xyz0, T0)
xyz1 = self.apply_transform(xyz1, T1)
else:
trans = np.identity(4)
# Voxelization and sampling
sel0 = ME.utils.sparse_quantize(xyz0 / self.voxel_size,
return_index=True)
sel1 = ME.utils.sparse_quantize(xyz1 / self.voxel_size,
return_index=True)
xyz0 = xyz0[np.random.choice(sel0, self.num_points)]
xyz1 = xyz1[np.random.choice(sel1, self.num_points)]
pointcloud1 = xyz0.T
pointcloud2 = xyz1.T
inv_trans = np.linalg.inv(trans)
R_ab = trans[:3, :3]
translation_ab = trans[:3, 3]
euler_ab = npmat2euler(np.expand_dims(R_ab, 0))
R_ba = inv_trans[:3, :3]
translation_ba = inv_trans[:3, 3]
euler_ba = npmat2euler(np.expand_dims(R_ba, 0))
# Selection
return pointcloud1.astype('float32'), pointcloud2.astype('float32'), R_ab.astype('float32'), \
translation_ab.astype('float32'), R_ba.astype('float32'), translation_ba.astype('float32'), \
euler_ab.astype('float32'), euler_ba.astype('float32')
def train_one_epoch(args, net, train_loader, opt):
net.train()
R_list = []
t_list = []
R_pred_list = []
t_pred_list = []
euler_list = []
for src, target, R, t, euler in tqdm(train_loader):
src = src.cuda()
target = target.cuda()
R = R.cuda()
t = t.cuda()
opt.zero_grad()
R_pred, t_pred, loss = net(src, target, R, t)
R_list.append(R.detach().cpu().numpy())
t_list.append(t.detach().cpu().numpy())
R_pred_list.append(R_pred.detach().cpu().numpy())
t_pred_list.append(t_pred.detach().cpu().numpy())
euler_list.append(euler.numpy())
loss.backward()
opt.step()
R = np.concatenate(R_list, axis=0)
t = np.concatenate(t_list, axis=0)
R_pred = np.concatenate(R_pred_list, axis=0)
t_pred = np.concatenate(t_pred_list, axis=0)
euler = np.concatenate(euler_list, axis=0)
euler_pred = npmat2euler(R_pred)
r_mse = np.mean((euler_pred - np.degrees(euler))**2)
r_rmse = np.sqrt(r_mse)
r_mae = np.mean(np.abs(euler_pred - np.degrees(euler)))
t_mse = np.mean((t - t_pred)**2)
t_rmse = np.sqrt(t_mse)
t_mae = np.mean(np.abs(t - t_pred))
return r_rmse, r_mae, t_rmse, t_mae
def test_one_epoch(args, net, test_loader):
net.eval()
R_list = []
t_list = []
R_pred_list = []
t_pred_list = []
euler_list = []
count = 0
for src, target, R, t, euler in tqdm(test_loader):
src = src.cuda()
target = target.cuda()
R_pred, t_pred, *_ = net(src, target)
src2 = torch.matmul(R_pred, src) + t_pred.unsqueeze(-1)
count += 1
R_list.append(R.numpy())
t_list.append(t.numpy())
R_pred_list.append(R_pred.detach().cpu().numpy())
t_pred_list.append(t_pred.detach().cpu().numpy())
euler_list.append(euler.numpy())
R = np.concatenate(R_list, axis=0)
t = np.concatenate(t_list, axis=0)
R_pred = np.concatenate(R_pred_list, axis=0)
t_pred = np.concatenate(t_pred_list, axis=0)
euler = np.concatenate(euler_list, axis=0)
euler_pred = npmat2euler(R_pred)
r_mse = np.mean((euler_pred - np.degrees(euler))**2)
r_rmse = np.sqrt(r_mse)
r_mae = np.mean(np.abs(euler_pred - np.degrees(euler)))
t_mse = np.mean((t - t_pred)**2)
t_rmse = np.sqrt(t_mse)
t_mae = np.mean(np.abs(t - t_pred))
return r_rmse, r_mae, t_rmse, t_mae
def test(args, net, test_loader, textio):
test_loss, test_rotations_ab, test_translations_ab, \
test_rotations_ab_pred, test_translations_ab_pred, \
test_eulers_ab = test_one_epoch(args, net, test_loader)
test_rotations_ab_pred_euler = npmat2euler(test_rotations_ab_pred)
test_r_mse_ab = np.mean(
(test_rotations_ab_pred_euler - np.degrees(test_eulers_ab))**2)
test_r_rmse_ab = np.sqrt(test_r_mse_ab)
test_r_mae_ab = np.mean(
np.abs(test_rotations_ab_pred_euler - np.degrees(test_eulers_ab)))
test_t_mse_ab = np.mean(
(test_translations_ab - test_translations_ab_pred)**2)
test_t_rmse_ab = np.sqrt(test_t_mse_ab)
test_t_mae_ab = np.mean(
np.abs(test_translations_ab - test_translations_ab_pred))
textio.cprint('==FINAL TEST==')
textio.cprint('A--------->B')
textio.cprint(
'EPOCH:: %d, Loss: %f, rot_MSE: %f, rot_RMSE: %f, '
'rot_MAE: %f, trans_MSE: %f, trans_RMSE: %f, trans_MAE: %f, Corr_Accuracy: %f'
% (-1, test_loss, test_r_mse_ab, test_r_rmse_ab, test_r_mae_ab,
test_t_mse_ab, test_t_rmse_ab, test_t_mae_ab, test_corr_accuracy))
def train(args, net, train_loader, test_loader, boardio, textio):
if args.use_sgd:
print("Use SGD")
opt = optim.SGD(net.parameters(),
lr=args.lr * 100,
momentum=args.momentum,
weight_decay=1e-4)
else:
print("Use Adam")
opt = optim.Adam(net.parameters(), lr=args.lr, weight_decay=1e-4)
scheduler = MultiStepLR(opt, milestones=[75, 150, 200], gamma=0.1)
best_test_loss = np.inf
best_test_cycle_loss = np.inf
best_test_mse_ab = np.inf
best_test_rmse_ab = np.inf
best_test_mae_ab = np.inf
best_test_r_mse_ab = np.inf
best_test_r_rmse_ab = np.inf
best_test_r_mae_ab = np.inf
best_test_t_mse_ab = np.inf
best_test_t_rmse_ab = np.inf
best_test_t_mae_ab = np.inf
best_test_mse_ba = np.inf
best_test_rmse_ba = np.inf
best_test_mae_ba = np.inf
best_test_r_mse_ba = np.inf
best_test_r_rmse_ba = np.inf
best_test_r_mae_ba = np.inf
best_test_t_mse_ba = np.inf
best_test_t_rmse_ba = np.inf
best_test_t_mae_ba = np.inf
for epoch in range(args.epochs):
scheduler.step()
train_loss, train_cycle_loss, \
train_mse_ab, train_mae_ab, train_mse_ba, train_mae_ba, train_rotations_ab, train_translations_ab, \
train_rotations_ab_pred, \
train_translations_ab_pred, train_rotations_ba, train_translations_ba, train_rotations_ba_pred, \
train_translations_ba_pred, train_eulers_ab, train_eulers_ba = train_one_epoch(args, net, train_loader, opt)
test_loss, test_cycle_loss, \
test_mse_ab, test_mae_ab, test_mse_ba, test_mae_ba, test_rotations_ab, test_translations_ab, \
test_rotations_ab_pred, \
test_translations_ab_pred, test_rotations_ba, test_translations_ba, test_rotations_ba_pred, \
test_translations_ba_pred, test_eulers_ab, test_eulers_ba = test_one_epoch(args, net, test_loader)
train_rmse_ab = np.sqrt(train_mse_ab)
test_rmse_ab = np.sqrt(test_mse_ab)
train_rmse_ba = np.sqrt(train_mse_ba)
test_rmse_ba = np.sqrt(test_mse_ba)
train_rotations_ab_pred_euler = npmat2euler(train_rotations_ab_pred)
train_r_mse_ab = np.mean(
(train_rotations_ab_pred_euler - np.degrees(train_eulers_ab))**2)
train_r_rmse_ab = np.sqrt(train_r_mse_ab)
train_r_mae_ab = np.mean(
np.abs(train_rotations_ab_pred_euler -
np.degrees(train_eulers_ab)))
train_t_mse_ab = np.mean(
(train_translations_ab - train_translations_ab_pred)**2)
train_t_rmse_ab = np.sqrt(train_t_mse_ab)
train_t_mae_ab = np.mean(
np.abs(train_translations_ab - train_translations_ab_pred))
train_rotations_ba_pred_euler = npmat2euler(train_rotations_ba_pred,
'xyz')
train_r_mse_ba = np.mean(
(train_rotations_ba_pred_euler - np.degrees(train_eulers_ba))**2)
train_r_rmse_ba = np.sqrt(train_r_mse_ba)
train_r_mae_ba = np.mean(
np.abs(train_rotations_ba_pred_euler -
np.degrees(train_eulers_ba)))
train_t_mse_ba = np.mean(
(train_translations_ba - train_translations_ba_pred)**2)
train_t_rmse_ba = np.sqrt(train_t_mse_ba)
train_t_mae_ba = np.mean(
np.abs(train_translations_ba - train_translations_ba_pred))
test_rotations_ab_pred_euler = npmat2euler(test_rotations_ab_pred)
test_r_mse_ab = np.mean(
(test_rotations_ab_pred_euler - np.degrees(test_eulers_ab))**2)
test_r_rmse_ab = np.sqrt(test_r_mse_ab)
test_r_mae_ab = np.mean(
np.abs(test_rotations_ab_pred_euler - np.degrees(test_eulers_ab)))
test_t_mse_ab = np.mean(
(test_translations_ab - test_translations_ab_pred)**2)
test_t_rmse_ab = np.sqrt(test_t_mse_ab)
test_t_mae_ab = np.mean(
np.abs(test_translations_ab - test_translations_ab_pred))
test_rotations_ba_pred_euler = npmat2euler(test_rotations_ba_pred,
'xyz')
test_r_mse_ba = np.mean(
(test_rotations_ba_pred_euler - np.degrees(test_eulers_ba))**2)
test_r_rmse_ba = np.sqrt(test_r_mse_ba)
test_r_mae_ba = np.mean(
np.abs(test_rotations_ba_pred_euler - np.degrees(test_eulers_ba)))
test_t_mse_ba = np.mean(
(test_translations_ba - test_translations_ba_pred)**2)
test_t_rmse_ba = np.sqrt(test_t_mse_ba)
test_t_mae_ba = np.mean(
np.abs(test_translations_ba - test_translations_ba_pred))
if best_test_loss >= test_loss:
best_test_loss = test_loss
best_test_cycle_loss = test_cycle_loss
best_test_mse_ab = test_mse_ab
best_test_rmse_ab = test_rmse_ab
best_test_mae_ab = test_mae_ab
best_test_r_mse_ab = test_r_mse_ab
best_test_r_rmse_ab = test_r_rmse_ab
best_test_r_mae_ab = test_r_mae_ab
best_test_t_mse_ab = test_t_mse_ab
best_test_t_rmse_ab = test_t_rmse_ab
best_test_t_mae_ab = test_t_mae_ab
best_test_mse_ba = test_mse_ba
best_test_rmse_ba = test_rmse_ba
best_test_mae_ba = test_mae_ba
best_test_r_mse_ba = test_r_mse_ba
best_test_r_rmse_ba = test_r_rmse_ba
best_test_r_mae_ba = test_r_mae_ba
best_test_t_mse_ba = test_t_mse_ba
best_test_t_rmse_ba = test_t_rmse_ba
best_test_t_mae_ba = test_t_mae_ba
if torch.cuda.device_count() > 1:
torch.save(
net.module.state_dict(),
'checkpoints/%s/models/model.best.t7' % args.exp_name)
else:
torch.save(
net.state_dict(),
'checkpoints/%s/models/model.best.t7' % args.exp_name)
textio.cprint('==TRAIN==')
textio.cprint('A--------->B')
textio.cprint(
'EPOCH:: %d, Loss: %f, Cycle Loss:, %f, MSE: %f, RMSE: %f, MAE: %f, rot_MSE: %f, rot_RMSE: %f, '
'rot_MAE: %f, trans_MSE: %f, trans_RMSE: %f, trans_MAE: %f' %
(epoch, train_loss, train_cycle_loss, train_mse_ab, train_rmse_ab,
train_mae_ab, train_r_mse_ab, train_r_rmse_ab, train_r_mae_ab,
train_t_mse_ab, train_t_rmse_ab, train_t_mae_ab))
textio.cprint('B--------->A')
textio.cprint(
'EPOCH:: %d, Loss: %f, MSE: %f, RMSE: %f, MAE: %f, rot_MSE: %f, rot_RMSE: %f, '
'rot_MAE: %f, trans_MSE: %f, trans_RMSE: %f, trans_MAE: %f' %
(epoch, train_loss, train_mse_ba, train_rmse_ba, train_mae_ba,
train_r_mse_ba, train_r_rmse_ba, train_r_mae_ba, train_t_mse_ba,
train_t_rmse_ba, train_t_mae_ba))
textio.cprint('==TEST==')
textio.cprint('A--------->B')
textio.cprint(
'EPOCH:: %d, Loss: %f, Cycle Loss: %f, MSE: %f, RMSE: %f, MAE: %f, rot_MSE: %f, rot_RMSE: %f, '
'rot_MAE: %f, trans_MSE: %f, trans_RMSE: %f, trans_MAE: %f' %
(epoch, test_loss, test_cycle_loss, test_mse_ab, test_rmse_ab,
test_mae_ab, test_r_mse_ab, test_r_rmse_ab, test_r_mae_ab,
test_t_mse_ab, test_t_rmse_ab, test_t_mae_ab))
textio.cprint('B--------->A')
textio.cprint(
'EPOCH:: %d, Loss: %f, MSE: %f, RMSE: %f, MAE: %f, rot_MSE: %f, rot_RMSE: %f, '
'rot_MAE: %f, trans_MSE: %f, trans_RMSE: %f, trans_MAE: %f' %
(epoch, test_loss, test_mse_ba, test_rmse_ba, test_mae_ba,
test_r_mse_ba, test_r_rmse_ba, test_r_mae_ba, test_t_mse_ba,
test_t_rmse_ba, test_t_mae_ba))
textio.cprint('==BEST TEST==')
textio.cprint('A--------->B')
textio.cprint(
'EPOCH:: %d, Loss: %f, Cycle Loss: %f, MSE: %f, RMSE: %f, MAE: %f, rot_MSE: %f, rot_RMSE: %f, '
'rot_MAE: %f, trans_MSE: %f, trans_RMSE: %f, trans_MAE: %f' %
(epoch, best_test_loss, best_test_cycle_loss, best_test_mse_ab,
best_test_rmse_ab, best_test_mae_ab, best_test_r_mse_ab,
best_test_r_rmse_ab, best_test_r_mae_ab, best_test_t_mse_ab,
best_test_t_rmse_ab, best_test_t_mae_ab))
textio.cprint('B--------->A')
textio.cprint(
'EPOCH:: %d, Loss: %f, MSE: %f, RMSE: %f, MAE: %f, rot_MSE: %f, rot_RMSE: %f, '
'rot_MAE: %f, trans_MSE: %f, trans_RMSE: %f, trans_MAE: %f' %
(epoch, best_test_loss, best_test_mse_ba, best_test_rmse_ba,
best_test_mae_ba, best_test_r_mse_ba, best_test_r_rmse_ba,
best_test_r_mae_ba, best_test_t_mse_ba, best_test_t_rmse_ba,
best_test_t_mae_ba))
boardio.add_scalar('A->B/train/loss', train_loss, epoch)
boardio.add_scalar('A->B/train/MSE', train_mse_ab, epoch)
boardio.add_scalar('A->B/train/RMSE', train_rmse_ab, epoch)
boardio.add_scalar('A->B/train/MAE', train_mae_ab, epoch)
boardio.add_scalar('A->B/train/rotation/MSE', train_r_mse_ab, epoch)
boardio.add_scalar('A->B/train/rotation/RMSE', train_r_rmse_ab, epoch)
boardio.add_scalar('A->B/train/rotation/MAE', train_r_mae_ab, epoch)
boardio.add_scalar('A->B/train/translation/MSE', train_t_mse_ab, epoch)
boardio.add_scalar('A->B/train/translation/RMSE', train_t_rmse_ab,
epoch)
boardio.add_scalar('A->B/train/translation/MAE', train_t_mae_ab, epoch)
boardio.add_scalar('B->A/train/loss', train_loss, epoch)
boardio.add_scalar('B->A/train/MSE', train_mse_ba, epoch)
boardio.add_scalar('B->A/train/RMSE', train_rmse_ba, epoch)
boardio.add_scalar('B->A/train/MAE', train_mae_ba, epoch)
boardio.add_scalar('B->A/train/rotation/MSE', train_r_mse_ba, epoch)
boardio.add_scalar('B->A/train/rotation/RMSE', train_r_rmse_ba, epoch)
boardio.add_scalar('B->A/train/rotation/MAE', train_r_mae_ba, epoch)
boardio.add_scalar('B->A/train/translation/MSE', train_t_mse_ba, epoch)
boardio.add_scalar('B->A/train/translation/RMSE', train_t_rmse_ba,
epoch)
boardio.add_scalar('B->A/train/translation/MAE', train_t_mae_ba, epoch)
############TEST
boardio.add_scalar('A->B/test/loss', test_loss, epoch)
boardio.add_scalar('A->B/test/MSE', test_mse_ab, epoch)
boardio.add_scalar('A->B/test/RMSE', test_rmse_ab, epoch)
boardio.add_scalar('A->B/test/MAE', test_mae_ab, epoch)
boardio.add_scalar('A->B/test/rotation/MSE', test_r_mse_ab, epoch)
boardio.add_scalar('A->B/test/rotation/RMSE', test_r_rmse_ab, epoch)
boardio.add_scalar('A->B/test/rotation/MAE', test_r_mae_ab, epoch)
boardio.add_scalar('A->B/test/translation/MSE', test_t_mse_ab, epoch)
boardio.add_scalar('A->B/test/translation/RMSE', test_t_rmse_ab, epoch)
boardio.add_scalar('A->B/test/translation/MAE', test_t_mae_ab, epoch)
boardio.add_scalar('B->A/test/loss', test_loss, epoch)
boardio.add_scalar('B->A/test/MSE', test_mse_ba, epoch)
boardio.add_scalar('B->A/test/RMSE', test_rmse_ba, epoch)
boardio.add_scalar('B->A/test/MAE', test_mae_ba, epoch)
boardio.add_scalar('B->A/test/rotation/MSE', test_r_mse_ba, epoch)
boardio.add_scalar('B->A/test/rotation/RMSE', test_r_rmse_ba, epoch)
boardio.add_scalar('B->A/test/rotation/MAE', test_r_mae_ba, epoch)
boardio.add_scalar('B->A/test/translation/MSE', test_t_mse_ba, epoch)
boardio.add_scalar('B->A/test/translation/RMSE', test_t_rmse_ba, epoch)
boardio.add_scalar('B->A/test/translation/MAE', test_t_mae_ba, epoch)
############BEST TEST
boardio.add_scalar('A->B/best_test/loss', best_test_loss, epoch)
boardio.add_scalar('A->B/best_test/MSE', best_test_mse_ab, epoch)
boardio.add_scalar('A->B/best_test/RMSE', best_test_rmse_ab, epoch)
boardio.add_scalar('A->B/best_test/MAE', best_test_mae_ab, epoch)
boardio.add_scalar('A->B/best_test/rotation/MSE', best_test_r_mse_ab,
epoch)
boardio.add_scalar('A->B/best_test/rotation/RMSE', best_test_r_rmse_ab,
epoch)
boardio.add_scalar('A->B/best_test/rotation/MAE', best_test_r_mae_ab,
epoch)
boardio.add_scalar('A->B/best_test/translation/MSE',
best_test_t_mse_ab, epoch)
boardio.add_scalar('A->B/best_test/translation/RMSE',
best_test_t_rmse_ab, epoch)
boardio.add_scalar('A->B/best_test/translation/MAE',
best_test_t_mae_ab, epoch)
boardio.add_scalar('B->A/best_test/loss', best_test_loss, epoch)
boardio.add_scalar('B->A/best_test/MSE', best_test_mse_ba, epoch)
boardio.add_scalar('B->A/best_test/RMSE', best_test_rmse_ba, epoch)
boardio.add_scalar('B->A/best_test/MAE', best_test_mae_ba, epoch)
boardio.add_scalar('B->A/best_test/rotation/MSE', best_test_r_mse_ba,
epoch)
boardio.add_scalar('B->A/best_test/rotation/RMSE', best_test_r_rmse_ba,
epoch)
boardio.add_scalar('B->A/best_test/rotation/MAE', best_test_r_mae_ba,
epoch)
boardio.add_scalar('B->A/best_test/translation/MSE',
best_test_t_mse_ba, epoch)
boardio.add_scalar('B->A/best_test/translation/RMSE',
best_test_t_rmse_ba, epoch)
boardio.add_scalar('B->A/best_test/translation/MAE',
best_test_t_mae_ba, epoch)
if torch.cuda.device_count() > 1:
torch.save(
net.module.state_dict(),
'checkpoints/%s/models/model.%d.t7' % (args.exp_name, epoch))
else:
torch.save(
net.state_dict(),
'checkpoints/%s/models/model.%d.t7' % (args.exp_name, epoch))
gc.collect()
def _train_one_epoch(self, epoch, train_loader, opt):
self.train()
total_loss = 0
rotations_ab = []
translations_ab = []
rotations_ab_pred = []
translations_ab_pred = []
eulers_ab = []
num_examples = 0
total_feature_alignment_loss = 0.0
total_cycle_consistency_loss = 0.0
total_scale_consensus_loss = 0.0
for data in tqdm(train_loader):
src, tgt, rotation_ab, translation_ab, rotation_ba, translation_ba, euler_ab, euler_ba = [d.cuda()
for d in data]
loss, feature_alignment_loss, cycle_consistency_loss, scale_consensus_loss,\
rotation_ab_pred, translation_ab_pred = self._train_one_batch(src, tgt, rotation_ab, translation_ab,
opt)
batch_size = src.size(0)
num_examples += batch_size
total_loss = total_loss + loss * batch_size
total_feature_alignment_loss = total_feature_alignment_loss + feature_alignment_loss * batch_size
total_cycle_consistency_loss = total_cycle_consistency_loss + cycle_consistency_loss * batch_size
total_scale_consensus_loss = total_scale_consensus_loss + scale_consensus_loss * batch_size
rotations_ab.append(rotation_ab.detach().cpu().numpy())
translations_ab.append(translation_ab.detach().cpu().numpy())
rotations_ab_pred.append(rotation_ab_pred.detach().cpu().numpy())
translations_ab_pred.append(translation_ab_pred.detach().cpu().numpy())
eulers_ab.append(euler_ab.cpu().numpy())
avg_loss = total_loss / num_examples
avg_feature_alignment_loss = total_feature_alignment_loss / num_examples
avg_cycle_consistency_loss = total_cycle_consistency_loss / num_examples
avg_scale_consensus_loss = total_scale_consensus_loss / num_examples
rotations_ab = np.concatenate(rotations_ab, axis=0)
translations_ab = np.concatenate(translations_ab, axis=0)
rotations_ab_pred = np.concatenate(rotations_ab_pred, axis=0)
translations_ab_pred = np.concatenate(translations_ab_pred, axis=0)
eulers_ab = np.degrees(np.concatenate(eulers_ab, axis=0))
eulers_ab_pred = npmat2euler(rotations_ab_pred)
r_ab_mse = np.mean((eulers_ab-eulers_ab_pred)**2)
r_ab_rmse = np.sqrt(r_ab_mse)
r_ab_mae = np.mean(np.abs(eulers_ab-eulers_ab_pred))
t_ab_mse = np.mean((translations_ab-translations_ab_pred)**2)
t_ab_rmse = np.sqrt(t_ab_mse)
t_ab_mae = np.mean(np.abs(translations_ab-translations_ab_pred))
r_ab_r2_score = r2_score(eulers_ab, eulers_ab_pred)
t_ab_r2_score = r2_score(translations_ab, translations_ab_pred)
info = {'arrow': 'A->B',
'epoch': epoch,
'stage': 'train',
'loss': avg_loss,
'feature_alignment_loss': avg_feature_alignment_loss,
'cycle_consistency_loss': avg_cycle_consistency_loss,
'scale_consensus_loss': avg_scale_consensus_loss,
'r_ab_mse': r_ab_mse,
'r_ab_rmse': r_ab_rmse,
'r_ab_mae': r_ab_mae,
't_ab_mse': t_ab_mse,
't_ab_rmse': t_ab_rmse,
't_ab_mae': t_ab_mae,
'r_ab_r2_score': r_ab_r2_score,
't_ab_r2_score': t_ab_r2_score}
self.logger.write(info)
return info
def train(args, model, train_loader, test_loader, boardio, textio, checkpoint):
learnable_params = filter(lambda p: p.requires_grad, model.parameters())
if args.optimizer == 'Adam':
optimizer = torch.optim.Adam(learnable_params)
else:
optimizer = torch.optim.SGD(learnable_params, lr=0.1)
if checkpoint is not None:
min_loss = checkpoint['min_loss']
optimizer.load_state_dict(checkpoint['optimizer'])
best_test_loss = np.inf
best_test_r_mse_ab = np.inf
best_test_r_rmse_ab = np.inf
best_test_r_mae_ab = np.inf
best_test_t_mse_ab = np.inf
best_test_t_rmse_ab = np.inf
best_test_t_mae_ab = np.inf
for epoch in range(args.start_epoch, args.epochs):
train_loss, train_rotations_ab, train_translations_ab, train_rotations_ab_pred, train_translations_ab_pred, train_eulers_ab = train_one_epoch(
args.device, model, train_loader, optimizer)
test_loss, test_rotations_ab, test_translations_ab, test_rotations_ab_pred, test_translations_ab_pred, test_eulers_ab = test_one_epoch(
args.device, model, test_loader)
train_rotations_ab_pred_euler = npmat2euler(train_rotations_ab_pred)
train_r_mse_ab = np.mean(
(train_rotations_ab_pred_euler - np.degrees(train_eulers_ab))**2)
train_r_rmse_ab = np.sqrt(train_r_mse_ab)
train_r_mae_ab = np.mean(
np.abs(train_rotations_ab_pred_euler -
np.degrees(train_eulers_ab)))
train_t_mse_ab = np.mean(
(train_translations_ab - train_translations_ab_pred)**2)
train_t_rmse_ab = np.sqrt(train_t_mse_ab)
train_t_mae_ab = np.mean(
np.abs(train_translations_ab - train_translations_ab_pred))
test_rotations_ab_pred_euler = npmat2euler(test_rotations_ab_pred)
test_r_mse_ab = np.mean(
(test_rotations_ab_pred_euler - np.degrees(test_eulers_ab))**2)
test_r_rmse_ab = np.sqrt(test_r_mse_ab)
test_r_mae_ab = np.mean(
np.abs(test_rotations_ab_pred_euler - np.degrees(test_eulers_ab)))
test_t_mse_ab = np.mean(
(test_translations_ab - test_translations_ab_pred)**2)
test_t_rmse_ab = np.sqrt(test_t_mse_ab)
test_t_mae_ab = np.mean(
np.abs(test_translations_ab - test_translations_ab_pred))
if test_loss < best_test_loss:
best_test_loss = test_loss
best_test_r_mse_ab = test_r_mse_ab
best_test_r_rmse_ab = test_r_rmse_ab
best_test_r_mae_ab = test_r_mae_ab
best_test_t_mse_ab = test_t_mse_ab
best_test_t_rmse_ab = test_t_rmse_ab
best_test_t_mae_ab = test_t_mae_ab
snap = {
'epoch': epoch + 1,
'model': model.state_dict(),
'min_loss': best_test_loss,
'optimizer': optimizer.state_dict(),
}
torch.save(
snap,
'checkpoints/%s/models/best_model_snap.t7' % (args.exp_name))
torch.save(model.state_dict(),
'checkpoints/%s/models/model.best.t7' % (args.exp_name))
torch.save(
model.feature_model.state_dict(),
'checkpoints/%s/models/best_ptnet_model.t7' % (args.exp_name))
torch.save(snap,
'checkpoints/%s/models/model_snap.t7' % (args.exp_name))
torch.save(model.state_dict(),
'checkpoints/%s/models/model.t7' % (args.exp_name))
torch.save(model.feature_model.state_dict(),
'checkpoints/%s/models/ptnet_model.t7' % (args.exp_name))
boardio.add_scalar('Train Loss', train_loss, epoch)
boardio.add_scalar('Test Loss', test_loss, epoch)
boardio.add_scalar('Best Test Loss', best_test_loss, epoch)
boardio.add_scalar('A->B/train/rotation/MSE', train_r_mse_ab, epoch)
boardio.add_scalar('A->B/train/rotation/RMSE', train_r_rmse_ab, epoch)
boardio.add_scalar('A->B/train/rotation/MAE', train_r_mae_ab, epoch)
boardio.add_scalar('A->B/train/translation/MSE', train_t_mse_ab, epoch)
boardio.add_scalar('A->B/train/translation/RMSE', train_t_rmse_ab,
epoch)
boardio.add_scalar('A->B/train/translation/MAE', train_t_mae_ab, epoch)
boardio.add_scalar('A->B/test/rotation/MSE', test_r_mse_ab, epoch)
boardio.add_scalar('A->B/test/rotation/RMSE', test_r_rmse_ab, epoch)
boardio.add_scalar('A->B/test/rotation/MAE', test_r_mae_ab, epoch)
boardio.add_scalar('A->B/test/translation/MSE', test_t_mse_ab, epoch)
boardio.add_scalar('A->B/test/translation/RMSE', test_t_rmse_ab, epoch)
boardio.add_scalar('A->B/test/translation/MAE', test_t_mae_ab, epoch)
boardio.add_scalar('A->B/best_test/rotation/MSE', best_test_r_mse_ab,
epoch)
boardio.add_scalar('A->B/best_test/rotation/RMSE', best_test_r_rmse_ab,
epoch)
boardio.add_scalar('A->B/best_test/rotation/MAE', best_test_r_mae_ab,
epoch)
boardio.add_scalar('A->B/best_test/translation/MSE',
best_test_t_mse_ab, epoch)
boardio.add_scalar('A->B/best_test/translation/RMSE',
best_test_t_rmse_ab, epoch)
boardio.add_scalar('A->B/best_test/translation/MAE',
best_test_t_mae_ab, epoch)
textio.cprint('==TRAIN==')
textio.cprint('A--------->B')
textio.cprint(
'EPOCH:: %d, Loss: %f, rot_MSE: %f, rot_RMSE: %f, '
'rot_MAE: %f, trans_MSE: %f, trans_RMSE: %f, trans_MAE: %f' %
(epoch, train_loss, train_r_mse_ab, train_r_rmse_ab,
train_r_mae_ab, train_t_mse_ab, train_t_rmse_ab, train_t_mae_ab))
textio.cprint('==TEST==')
textio.cprint('A--------->B')
textio.cprint(
'EPOCH:: %d, Loss: %f, rot_MSE: %f, rot_RMSE: %f, '
'rot_MAE: %f, trans_MSE: %f, trans_RMSE: %f, trans_MAE: %f' %
(epoch, test_loss, test_r_mse_ab, test_r_rmse_ab, test_r_mae_ab,
test_t_mse_ab, test_t_rmse_ab, test_t_mae_ab))
textio.cprint('==BEST TEST==')
textio.cprint('A--------->B')
textio.cprint(
'EPOCH:: %d, Loss: %f, rot_MSE: %f, rot_RMSE: %f, '
'rot_MAE: %f, trans_MSE: %f, trans_RMSE: %f, trans_MAE: %f' %
(epoch, best_test_loss, best_test_r_mse_ab, best_test_r_rmse_ab,
best_test_r_mae_ab, best_test_t_mse_ab, best_test_t_rmse_ab,
best_test_t_mae_ab))
