main.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-


from __future__ import print_function
import os
import gc
import argparse
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.optim.lr_scheduler import MultiStepLR
from data import CustomDataset
from model import DCP, knn
from util import transform_point_cloud, npmat2euler, dump_point_cloud
import numpy as np
from torch.utils.data import DataLoader
from tensorboardX import SummaryWriter
from tqdm import tqdm
import trimesh
import json


# Part of the code is referred from: https://github.com/floodsung/LearningToCompare_FSL

class IOStream:
    def __init__(self, path):
        self.f = open(path, 'a')

    def cprint(self, text):
        print(text)
        self.f.write(text + '\n')
        self.f.flush()

    def close(self):
        self.f.close()


def _init_(args):
    if not os.path.exists('checkpoints'):
        os.makedirs('checkpoints')
    if not os.path.exists('checkpoints/' + args.exp_name):
        os.makedirs('checkpoints/' + args.exp_name)
    if not os.path.exists('checkpoints/' + args.exp_name + '/' + 'models'):
        os.makedirs('checkpoints/' + args.exp_name + '/' + 'models')
    os.system('cp main.py checkpoints' + '/' + args.exp_name + '/' + 'main.py.backup')
    os.system('cp model.py checkpoints' + '/' + args.exp_name + '/' + 'model.py.backup')
    os.system('cp data.py checkpoints' + '/' + args.exp_name + '/' + 'data.py.backup')


def test_one_epoch(args, net, test_loader):
    net.eval()
    mse_ab = 0
    mae_ab = 0
    mse_ba = 0
    mae_ba = 0

    total_loss = 0
    total_cycle_loss = 0
    num_examples = 0
    rotations_ab = []
    translations_ab = []
    rotations_ab_pred = []
    translations_ab_pred = []

    rotations_ba = []
    translations_ba = []
    rotations_ba_pred = []
    translations_ba_pred = []

    eulers_ab = []
    eulers_ba = []
    
    for idx, (src, target, rotation_ab, translation_ab, rotation_ba, translation_ba, euler_ab, euler_ba, \
            color_src, color_target) in enumerate(tqdm(test_loader)):

        src = src.to(args.device)
        target = target.to(args.device)
        color_src = color_src.to(args.device)
        color_target = color_target.to(args.device)

        rotation_ab = rotation_ab.to(args.device)
        translation_ab = translation_ab.to(args.device)
        rotation_ba = rotation_ba.to(args.device)
        translation_ba = translation_ba.to(args.device)

        batch_size = src.size(0)
        num_examples += batch_size
        rotation_ab_pred, translation_ab_pred, rotation_ba_pred, translation_ba_pred,assignment_ab = \
            net(src, target, color_src, color_target)

        # save rotation and translation
        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.numpy())
        ##
        rotations_ba.append(rotation_ba.detach().cpu().numpy())
        translations_ba.append(translation_ba.detach().cpu().numpy())
        rotations_ba_pred.append(rotation_ba_pred.detach().cpu().numpy())
        translations_ba_pred.append(translation_ba_pred.detach().cpu().numpy())
        eulers_ba.append(euler_ba.numpy())

        transformed_src = transform_point_cloud(src, rotation_ab_pred, translation_ab_pred)

        transformed_target = transform_point_cloud(target, rotation_ba_pred, translation_ba_pred)

        for i in range(batch_size):
            dump_point_cloud(f'results/{args.exp_name}/outputs/{idx * batch_size + i}_source', src[i], color_src[i])
            dump_point_cloud(f'results/{args.exp_name}/outputs/{idx * batch_size + i}_target', target[i], color_target[i])
            dump_point_cloud(f'results/{args.exp_name}/outputs/{idx * batch_size + i}_transformed_source', transformed_src[i], color_src[i])
            dump_point_cloud(f'results/{args.exp_name}/outputs/{idx * batch_size + i}_transformed_target', transformed_target[i], color_target[i])

        ###########################
        identity = torch.eye(3).to(args.device).unsqueeze(0).repeat(batch_size, 1, 1)
        loss = F.mse_loss(torch.matmul(rotation_ab_pred.transpose(2, 1), rotation_ab), identity) \
            + F.mse_loss(translation_ab_pred, translation_ab)
        if args.cycle:
            rotation_loss = F.mse_loss(torch.matmul(rotation_ba_pred, rotation_ab_pred), identity.clone())
            translation_loss = torch.mean((torch.matmul(rotation_ba_pred.transpose(2, 1),
                                                        translation_ab_pred.view(batch_size, 3, 1)).view(batch_size, 3)
                                           + translation_ba_pred) ** 2, dim=[0, 1])
            cycle_loss = rotation_loss + translation_loss

            loss = loss + cycle_loss * 0.1

        total_loss += loss.item() * batch_size

        if args.cycle:
            total_cycle_loss = total_cycle_loss + cycle_loss.item() * 0.1 * batch_size

        mse_ab += torch.mean((transformed_src - target) ** 2, dim=[0, 1, 2]).item() * batch_size
        mae_ab += torch.mean(torch.abs(transformed_src - target), dim=[0, 1, 2]).item() * batch_size

        mse_ba += torch.mean((transformed_target - src) ** 2, dim=[0, 1, 2]).item() * batch_size
        mae_ba += torch.mean(torch.abs(transformed_target - src), dim=[0, 1, 2]).item() * batch_size

        # m.export("pooints.ply","ply")
        # m = trimesh.PointCloud(vertices=verts,vertex_colors=colors)
        # print(src.cpu()[0].shape)
        # print("Yodelei")
        #m_src = trimesh.PointCloud(vertices=src.cpu().detach().numpy()[0].T, vertex_colors=color_src.cpu().detach().numpy()[0].T)
        #m_tgt = trimesh.PointCloud(vertices=target.cpu().detach().numpy()[0].T, vertex_colors=color_target.cpu().detach().numpy()[0].T)
        #m_src_trans = trimesh.PointCloud(vertices=transformed_src.cpu().detach().numpy()[0].T, vertex_colors=color_src.cpu().detach().numpy()[0].T)
        #m_tgt_trans = trimesh.PointCloud(vertices=transformed_target.cpu().detach().numpy()[0].T, vertex_colors=color_target.cpu().detach().numpy()[0].T)
        # m_src.export("src.ply","ply")
        # m_tgt.export("tgt.ply","ply")
        # m_src_trans.export("src_trans.ply","ply")
        # m_tgt_trans.export("tgt_trans.ply","ply")

    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)

    rotations_ba = np.concatenate(rotations_ba, axis=0)
    translations_ba = np.concatenate(translations_ba, axis=0)
    rotations_ba_pred = np.concatenate(rotations_ba_pred, axis=0)
    translations_ba_pred = np.concatenate(translations_ba_pred, axis=0)

    eulers_ab = np.concatenate(eulers_ab, axis=0)
    eulers_ba = np.concatenate(eulers_ba, axis=0)

    return total_loss * 1.0 / num_examples, total_cycle_loss / num_examples, \
        mse_ab * 1.0 / num_examples, mae_ab * 1.0 / num_examples, \
        mse_ba * 1.0 / num_examples, mae_ba * 1.0 / num_examples, rotations_ab, \
        translations_ab, rotations_ab_pred, translations_ab_pred, rotations_ba, \
        translations_ba, rotations_ba_pred, translations_ba_pred, eulers_ab, eulers_ba


def arap_energy(src, trgt, R,assignment_ab):
    smm = 0
    neighs = knn(src, 5)
    for b in range(src.shape[0]):  # batch_size
        for i in range(src[b].shape[0]):
            for j in range(5):
                n_ind_s = neighs[b, i, j]
                n_ind_t = assignment_ab[b,n_ind_s]
                smm += torch.linalg.norm((torch.matmul(R[b].T, trgt[b, :,assignment_ab[b,i]] - trgt[b,:, n_ind_t])) - (src[b,:, i] - src[b,:, n_ind_s]))

    return smm
def train_one_epoch(args, net, train_loader, opt):
    net.train()

    mse_ab = 0
    mae_ab = 0
    mse_ba = 0
    mae_ba = 0

    total_loss = 0
    total_cycle_loss = 0
    num_examples = 0
    tot_ar_loss= 0
    rotations_ab = []
    translations_ab = []
    rotations_ab_pred = []
    translations_ab_pred = []

    rotations_ba = []
    translations_ba = []
    rotations_ba_pred = []
    translations_ba_pred = []

    eulers_ab = []
    eulers_ba = []

    for src, target, rotation_ab, translation_ab, rotation_ba, translation_ba, euler_ab, euler_ba, \
            color_src, color_target in tqdm(train_loader):
        src = src.to(args.device)
        target = target.to(args.device)
        color_src = color_src.to(args.device)
        color_target = color_target.to(args.device)
        rotation_ab = rotation_ab.to(args.device)
        translation_ab = translation_ab.to(args.device)
        rotation_ba = rotation_ba.to(args.device)
        translation_ba = translation_ba.to(args.device)

        batch_size = src.size(0)
        opt.zero_grad()
        num_examples += batch_size
        rotation_ab_pred, translation_ab_pred, rotation_ba_pred, translation_ba_pred,assignment_ab = \
            net(src, target, color_src, color_target)
        # save rotation and translation
        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.numpy())
        ##
        rotations_ba.append(rotation_ba.detach().cpu().numpy())
        translations_ba.append(translation_ba.detach().cpu().numpy())
        rotations_ba_pred.append(rotation_ba_pred.detach().cpu().numpy())
        translations_ba_pred.append(translation_ba_pred.detach().cpu().numpy())
        eulers_ba.append(euler_ba.numpy())

        transformed_src = transform_point_cloud(src, rotation_ab_pred, translation_ab_pred)

        transformed_target = transform_point_cloud(target, rotation_ba_pred, translation_ba_pred)
        ###########################
        identity = torch.eye(3).to(args.device).unsqueeze(0).repeat(batch_size, 1, 1)
        loss = F.mse_loss(torch.matmul(rotation_ab_pred.transpose(2, 1), rotation_ab), identity) \
            + F.mse_loss(translation_ab_pred, translation_ab)
        if args.cycle:
            rotation_loss = F.mse_loss(torch.matmul(rotation_ba_pred, rotation_ab_pred), identity.clone())
            translation_loss = torch.mean((torch.matmul(rotation_ba_pred.transpose(2, 1),
                                                        translation_ab_pred.view(batch_size, 3, 1)).view(batch_size, 3)
                                           + translation_ba_pred) ** 2, dim=[0, 1])
            cycle_loss = rotation_loss + translation_loss

            loss = loss + cycle_loss * 0.1
        if args.arap and assignment_ab is not None:
            ar = arap_energy(src,target,rotation_ab,assignment_ab)
            loss += 0.2*ar # lambda hyperparameter
            tot_ar_loss = tot_ar_loss + ar

        loss.backward()
        opt.step()
        total_loss += loss.item() * batch_size

        if args.cycle:
            total_cycle_loss = total_cycle_loss + cycle_loss.item() * 0.1 * batch_size

        mse_ab += torch.mean((transformed_src - target) ** 2, dim=[0, 1, 2]).item() * batch_size
        mae_ab += torch.mean(torch.abs(transformed_src - target), dim=[0, 1, 2]).item() * batch_size

        mse_ba += torch.mean((transformed_target - src) ** 2, dim=[0, 1, 2]).item() * batch_size
        mae_ba += torch.mean(torch.abs(transformed_target - src), dim=[0, 1, 2]).item() * batch_size
        avg_ar = tot_ar_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)

    rotations_ba = np.concatenate(rotations_ba, axis=0)
    translations_ba = np.concatenate(translations_ba, axis=0)
    rotations_ba_pred = np.concatenate(rotations_ba_pred, axis=0)
    translations_ba_pred = np.concatenate(translations_ba_pred, axis=0)

    eulers_ab = np.concatenate(eulers_ab, axis=0)
    eulers_ba = np.concatenate(eulers_ba, axis=0)

    return total_loss * 1.0 / num_examples, total_cycle_loss / num_examples, \
        mse_ab * 1.0 / num_examples, mae_ab * 1.0 / num_examples, \
        mse_ba * 1.0 / num_examples, mae_ba * 1.0 / num_examples, rotations_ab, \
        translations_ab, rotations_ab_pred, translations_ab_pred, rotations_ba, \
        translations_ba, rotations_ba_pred, translations_ba_pred, eulers_ab, eulers_ba,avg_ar


def test(args, net, test_loader, boardio, textio):

    result_dir = f'results/{args.exp_name}'
    output_dir = f'{result_dir}/outputs'
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    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)

    scores = {}
    test_rotations_ab_pred_euler = npmat2euler(test_rotations_ab_pred)
    scores['test_r_mse_ab'] = np.mean((test_rotations_ab_pred_euler - np.degrees(test_eulers_ab)) ** 2)
    scores['test_r_rmse_ab'] = np.sqrt(scores['test_r_mse_ab'])
    scores['test_r_mae_ab'] = np.mean(np.abs(test_rotations_ab_pred_euler - np.degrees(test_eulers_ab)))
    scores['test_t_mse_ab'] = np.mean((test_translations_ab - test_translations_ab_pred) ** 2)
    scores['test_t_rmse_ab'] = np.sqrt(scores['test_t_mse_ab'])
    scores['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')
    scores['test_r_mse_ba'] = np.mean((test_rotations_ba_pred_euler - np.degrees(test_eulers_ba)) ** 2)
    scores['test_r_rmse_ba'] = np.sqrt(scores['test_r_mse_ba'])
    scores['test_r_mae_ba'] = np.mean(np.abs(test_rotations_ba_pred_euler - np.degrees(test_eulers_ba)))
    scores['test_t_mse_ba'] = np.mean((test_translations_ba - test_translations_ba_pred) ** 2)
    scores['test_t_rmse_ba'] = np.sqrt(scores['test_t_mse_ba'])
    scores['test_t_mae_ba'] = np.mean(np.abs(test_translations_ba - test_translations_ba_pred))

    with open(f"{result_dir}/scores.json", 'w') as fout:
        json_dumps_str = json.dumps(str(scores), indent=4)
        print(json_dumps_str, file=fout)

    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,
                     scores['test_r_mse_ab'], scores['test_r_rmse_ab'],
                     scores['test_r_mae_ab'], scores['test_t_mse_ab'], scores['test_t_rmse_ab'], scores['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, scores['test_r_mse_ba'], scores['test_r_rmse_ba'],
                     scores['test_r_mae_ba'], scores['test_t_mse_ba'], scores['test_t_rmse_ba'], scores['test_t_mae_ba']))


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):
        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,arap_loss = 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.is_available() and 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))

        boardio.add_scalar('ARAP LOSS', arap_loss, epoch)
        textio.cprint('AVG ARAP LOSS: %f' % arap_loss)

        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.is_available() and 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()
        scheduler.step()


def main():
    parser = argparse.ArgumentParser(description='Point Cloud Registration')
    parser.add_argument('--exp-name', type=str, default='exp', metavar='N',
                        help='Name of the experiment')
    parser.add_argument('--model', type=str, default='dcp', metavar='N',
                        choices=['dcp'],
                        help='Model to use, [dcp]')
    parser.add_argument('--emb-nn', type=str, default='pointnet', metavar='N',
                        choices=['pointnet', 'dgcnn'],
                        help='Embedding nn to use, [pointnet, dgcnn]')
    parser.add_argument('--pointer', type=str, default='transformer', metavar='N',
                        choices=['identity', 'transformer'],
                        help='Attention-based pointer generator to use, [identity, transformer]')
    parser.add_argument('--head', type=str, default='svd', metavar='N',
                        choices=['mlp', 'svd', ],
                        help='Head to use, [mlp, svd]')
    parser.add_argument('--emb_dims', type=int, default=512, metavar='N',
                        help='Dimension of embeddings')
    parser.add_argument('--n_blocks', type=int, default=1, metavar='N',
                        help='Num of blocks of encoder&decoder')
    parser.add_argument('--n_heads', type=int, default=4, metavar='N',
                        help='Num of heads in multiheadedattention')
    parser.add_argument('--ff_dims', type=int, default=1024, metavar='N',
                        help='Num of dimensions of fc in transformer')
    parser.add_argument('--dropout', type=float, default=0.0, metavar='N',
                        help='Dropout ratio in transformer')
    parser.add_argument('--batch-size', type=int, default=32, metavar='batch_size',
                        help='Size of batch)')
    parser.add_argument('--test-batch-size', type=int, default=10, metavar='batch_size',
                        help='Size of batch)')
    parser.add_argument('--epochs', type=int, default=250, metavar='N',
                        help='number of episode to train ')
    parser.add_argument('--use-sgd', action='store_true', default=False,
                        help='Use SGD')
    parser.add_argument('--lr', type=float, default=0.001, metavar='LR',
                        help='learning rate (default: 0.001, 0.1 if using sgd)')
    parser.add_argument('--momentum', type=float, default=0.9, metavar='M',
                        help='SGD momentum (default: 0.9)')
    parser.add_argument('--no-cuda', action='store_true', default=False,
                        help='enables CUDA training')
    parser.add_argument('--seed', type=int, default=1234, metavar='S',
                        help='random seed (default: 1)')
    parser.add_argument('--eval', action='store_true', default=False,
                        help='evaluate the model')
    parser.add_argument('--cycle', type=bool, default=False, metavar='N',
                        help='Whether to use cycle consistency')
    parser.add_argument('--gaussian-noise', type=bool, default=False, metavar='N',
                        help='Wheter to add gaussian noise')
    parser.add_argument('--unseen', type=bool, default=False, metavar='N',
                        help='Wheter to test on unseen category')
    parser.add_argument('--num-points', type=int, default=1024, metavar='N',
                        help='Num of points to use')
    parser.add_argument('--dataset', type=str, default='modelnet40', choices=['modelnet40', 'mixamo','tumrgbd'], metavar='N',
                        help='dataset to use')
    parser.add_argument('--factor', type=float, default=4, metavar='N',
                        help='Divided factor for rotations')
    parser.add_argument('--model-path', type=str, default='', metavar='N',
                        help='Pretrained model path')
    parser.add_argument('--matching-method', type=str, default='softmax', metavar='N', choices=['softmax', 'sink_horn'],
                        help='The point matching method')
    parser.add_argument('--device', type=str, default='cuda', metavar='N',
                        help='Pretrained model path')
    parser.add_argument('--no_slack', action='store_true', help='If set, will not have a slack column.')
    parser.add_argument('--num_sk_iter', type=int, default=5,
                        help='Number of inner iterations used in sinkhorn normalization')
    parser.add_argument('--use-color', type=bool, default=False, metavar='N',
                        help='Flag for using the color as input')
    parser.add_argument('--different-sampling', type=bool, default=False, metavar='N',
                        help='Flag for using different sampling in source and target')
    parser.add_argument('--arap', type=bool, default=False, metavar='N',
                        help='Flag for ARAP regularizer')

    args = parser.parse_args()
    torch.manual_seed(args.seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(args.seed)
        torch.backends.cudnn.deterministic = True
    np.random.seed(args.seed)

    boardio = SummaryWriter(log_dir='checkpoints/' + args.exp_name)

    # Read Device from run arguments and initilize the corresponding device
    if 'cuda' in args.device and torch.cuda.is_available():
        args.device = torch.device(args.device)
    elif 'cpu' in args.device:
        args.device = torch.device(args.device)
    else:
        raise NotImplementedError(f'Device:"{args.device}" is not implemented or not visible')
    _init_(args)

    textio = IOStream('checkpoints/' + args.exp_name + '/run.log')
    textio.cprint(str(args))

    if args.dataset in ['modelnet40', 'mixamo','tumrgbd']:
        train_loader = DataLoader(
            CustomDataset(num_points=args.num_points, partition='train', gaussian_noise=args.gaussian_noise,
                          unseen=args.unseen, factor=args.factor, dataset=args.dataset, use_color=args.use_color,
                          different_sampling=args.different_sampling),
            batch_size=args.batch_size, shuffle=True, drop_last=True)
        test_loader = DataLoader(
            CustomDataset(num_points=args.num_points, partition='test', gaussian_noise=args.gaussian_noise,
                          unseen=args.unseen, factor=args.factor, dataset=args.dataset, use_color=args.use_color,
                          different_sampling=args.different_sampling),
            batch_size=args.test_batch_size, shuffle=False, drop_last=False)
    else:
        raise Exception("not implemented")

    if args.model == 'dcp':
        net = DCP(args).to(args.device)
        if args.eval:
            if args.model_path == '':
                model_path = 'checkpoints' + '/' + args.exp_name + '/models/model.best.t7'
            else:
                model_path = args.model_path
                print(model_path)
            if not os.path.exists(model_path):
                print("can't find pretrained model")
                return
            net.load_state_dict(torch.load(model_path, map_location=torch.device(args.device)), strict=False)
        if torch.cuda.is_available() and torch.cuda.device_count() > 1:
            net = nn.DataParallel(net)
            print("Let's use", torch.cuda.device_count(), "GPUs!")
    else:
        raise Exception('Not implemented')
    if args.eval:
        test(args, net, test_loader, boardio, textio)
    else:
        train(args, net, train_loader, test_loader, boardio, textio)

    print('FINISH')
    boardio.close()


if __name__ == '__main__':
    main()