def __init__(self):
model = get_model()
print('local_rank', args.node_rank)
torch.cuda.set_device(args.node_rank)
model = model.cuda()
model = torch.nn.parallel.DistributedDataParallel(
model, device_ids=[args.node_rank], output_device=args.node_rank)
print("From rank %d: start training, time:%s" %
(args.node_rank, time.strftime("%Y-%m-%d %H:%M:%S")))
criterion = get_loss
optimizer = get_optimizer(model)
self.ce_loss = torch.nn.BCELoss()
self.mse_loss = torch.nn.MSELoss()
self.wing_loss = get_wing_loss
self.focal_loss = get_focal_loss
super(FaceAlign_Trainer, self).__init__(model=model,
criterion=criterion,
optimizer=optimizer,
opt=opt)
from tensorboardX import SummaryWriter
self.writer = SummaryWriter(opt.vis_dir)
self.best_metric = 10000
for m in all_metrics:
self.metric_meter[m] = meter.AverageValueMeter()
def __init__(self, model):
criterion = get_loss
optimizer = get_optimizer(model)
self.ce_loss = torch.nn.BCELoss()
self.mse_loss = torch.nn.MSELoss()
self.wing_loss = get_wing_loss
self.focal_loss = get_focal_loss
self.n_iter = 0
self.n_plot = 0
super(FaceAlign_Trainer, self).__init__(model=model, criterion=criterion, optimizer=optimizer, opt=opt)
if args.rank == 0:
from tensorboardX import SummaryWriter
self.writer = SummaryWriter(opt.vis_dir)
self.best_metric = 10000
for m in all_metrics:
self.metric_meter[m] = meter.AverageValueMeter()
def __call__(self, network, P1, LABEL1, P2, LABEL2, local_fix=None, latent=False, latent_P1=None, latent_P2=None,
distChamfer=None):
"""
:param network:
:param P1: input points 1 -> batch, num_points, 3
:param P2: input points 2 -> batch, num_points, 3
:param local_fix: Vector storing constants for batch_idx function
:param latent: Boolean to decide if latent vector should be stored for slight optimisation of runtime
:param latent_P1: Latent vector 1 precomputed
:param latent_P2: Latent vector 2 precomputed
:param distChamfer: Chamfer distance function
:return: P2 reconstructed from P1, latent vectors, chamfer outputs
"""
if latent:
P2_P1, latent_vector_P1, latent_vector_P2 = network.forward_classic_with_latent(
P1.transpose(2, 1).contiguous(), P2.transpose(2, 1).contiguous(), x_latent=latent_P1,
y_latent=latent_P2) # forward pass
else:
P2_P1 = network(P1.transpose(2, 1).contiguous(), P2.transpose(2, 1).contiguous()) # forward pass
P2_P1 = P2_P1.transpose(2, 1).contiguous()
# reconstruction losses
dist1, dist2, idx1, idx2 = distChamfer(P2_P1, P2)
loss = meter.AverageValueMeter()
for shape in range(P1.size(0)):
for label in set(LABEL1[shape].cpu().numpy()):
# A for loop on the batch_size in neccessary as shape parts don't have the same number of points for each shape.
try:
dist1_label, dist2_label, idx1_label, idx2_label = distChamfer(
P2_P1[shape][LABEL1[shape] == label].unsqueeze(0),
P2[shape][LABEL2[shape] == label].unsqueeze(0))
loss.update(self.chamfer_loss.forward(dist1_label, dist2_label))
except:
continue
# TODO :find a good way to make sure points with no matching labels are not accounted for in cycle consistency
# This is the critical step that allows the batched indexing of the computation of the cycle losses to runs smoothly.
if not local_fix is None:
idx1 = batch_idx(idx1, local_fix)
idx2 = batch_idx(idx2, local_fix)
if latent:
return P2_P1, dist1, dist2, idx1.long(), idx2.long(), latent_vector_P1, latent_vector_P2, loss.avg
else:
return P2_P1, dist1, dist2, idx1.long(), idx2.long(), loss.avg
import numpy as np
import open3d
import torch
import copy
import time
import meter
timings_ICP = meter.AverageValueMeter() # initialize iou for this shape
def ICP(source, target):
"""
:param source: source point cloud
:param target: target point cloud
:return: source pointcloud registered
"""
start = time.time()
use_torch = False
if isinstance(source, torch.Tensor):
use_torch = True
source = source.squeeze().cpu().numpy()
if isinstance(target, torch.Tensor):
target = target.squeeze().cpu().numpy()
pcd_target = open3d.PointCloud()
pcd_target.points = open3d.Vector3dVector(target)
pcd_source = open3d.PointCloud()
pcd_source.points = open3d.Vector3dVector(source)