def eval_pcpnet(opt):
opt.models = opt.models.split()
if opt.seed < 0:
opt.seed = random.randint(1, 10000)
device = torch.device("cpu" if opt.gpu_idx < 0 else "cuda:%d" % opt.gpu_idx)
for model_name in opt.models:
print("Random Seed: %d" % (opt.seed))
random.seed(opt.seed)
torch.manual_seed(opt.seed)
model_filename = os.path.join(opt.modeldir, model_name+opt.modelpostfix)
param_filename = os.path.join(opt.modeldir, model_name+opt.parmpostfix)
# load model and training parameters
trainopt = torch.load(param_filename)
if opt.batchSize == 0:
model_batchSize = trainopt.batchSize
else:
model_batchSize = opt.batchSize
# get indices in targets and predictions corresponding to each output
pred_dim = 0
output_pred_ind = []
for o in trainopt.outputs:
if o == 'unoriented_normals' or o == 'oriented_normals':
output_pred_ind.append(pred_dim)
pred_dim += 3
elif o == 'max_curvature' or o == 'min_curvature':
output_pred_ind.append(pred_dim)
pred_dim += 1
else:
raise ValueError('Unknown output: %s' % (o))
dataset = PointcloudPatchDataset(
root=opt.indir, shape_list_filename=opt.dataset,
patch_radius=trainopt.patch_radius,
points_per_patch=trainopt.points_per_patch,
patch_features=[],
seed=opt.seed,
use_pca=trainopt.use_pca,
center=trainopt.patch_center,
point_tuple=trainopt.point_tuple,
sparse_patches=opt.sparse_patches,
cache_capacity=opt.cache_capacity)
if opt.sampling == 'full':
datasampler = SequentialPointcloudPatchSampler(dataset)
elif opt.sampling == 'sequential_shapes_random_patches':
datasampler = SequentialShapeRandomPointcloudPatchSampler(
dataset,
patches_per_shape=opt.patches_per_shape,
seed=opt.seed,
sequential_shapes=True,
identical_epochs=False)
else:
raise ValueError('Unknown sampling strategy: %s' % opt.sampling)
dataloader = torch.utils.data.DataLoader(
dataset,
sampler=datasampler,
batch_size=model_batchSize,
num_workers=int(opt.workers))
if len(trainopt.patch_radius) == 1:
regressor = PCPNet(
num_points=trainopt.points_per_patch,
output_dim=pred_dim,
use_point_stn=trainopt.use_point_stn,
use_feat_stn=trainopt.use_feat_stn,
sym_op=trainopt.sym_op,
point_tuple=trainopt.point_tuple)
else:
regressor = MSPCPNet(
num_scales=len(trainopt.patch_radius),
num_points=trainopt.points_per_patch,
output_dim=pred_dim,
use_point_stn=trainopt.use_point_stn,
use_feat_stn=trainopt.use_feat_stn,
sym_op=trainopt.sym_op,
point_tuple=trainopt.point_tuple)
regressor.load_state_dict(torch.load(model_filename))
regressor.to(device)
regressor.eval()
shape_ind = 0
shape_patch_offset = 0
if opt.sampling == 'full':
shape_patch_count = dataset.shape_patch_count[shape_ind]
elif opt.sampling == 'sequential_shapes_random_patches':
shape_patch_count = min(opt.patches_per_shape, dataset.shape_patch_count[shape_ind])
else:
raise ValueError('Unknown sampling strategy: %s' % opt.sampling)
shape_properties = torch.zeros(shape_patch_count, pred_dim, dtype=torch.float, device=device)
# append model name to output directory and create directory if necessary
model_outdir = os.path.join(opt.outdir, model_name)
if not os.path.exists(model_outdir):
os.makedirs(model_outdir)
num_batch = len(dataloader)
batch_enum = enumerate(dataloader, 0)
for batchind, data in batch_enum:
# get batch and upload to GPU
points, data_trans = data
points = points.transpose(2, 1)
points = points.to(device)
data_trans = data_trans.to(device)
with torch.no_grad():
pred, trans, _, _ = regressor(points)
# post-processing of the prediction
for oi, o in enumerate(trainopt.outputs):
if o == 'unoriented_normals' or o == 'oriented_normals':
o_pred = pred[:, output_pred_ind[oi]:output_pred_ind[oi]+3]
if trainopt.use_point_stn:
# transform predictions with inverse transform
# since we know the transform to be a rotation (QSTN), the transpose is the inverse
o_pred[:, :] = torch.bmm(o_pred.unsqueeze(1), trans.transpose(2, 1)).squeeze(dim=1)
if trainopt.use_pca:
# transform predictions with inverse pca rotation (back to world space)
o_pred[:, :] = torch.bmm(o_pred.unsqueeze(1), data_trans.transpose(2, 1)).squeeze(dim=1)
# normalize normals
o_pred_len = torch.max(o_pred.new_tensor([sys.float_info.epsilon*100]), o_pred.norm(p=2, dim=1, keepdim=True))
o_pred = o_pred / o_pred_len
elif o == 'max_curvature' or o == 'min_curvature':
o_pred = pred[:, output_pred_ind[oi]:output_pred_ind[oi]+1]
# undo patch size normalization:
o_pred[:, :] = o_pred / dataset.patch_radius_absolute[shape_ind][0]
else:
raise ValueError('Unsupported output type: %s' % (o))
print('[%s %d/%d] shape %s' % (model_name, batchind, num_batch-1, dataset.shape_names[shape_ind]))
batch_offset = 0
while batch_offset < pred.size(0):
shape_patches_remaining = shape_patch_count-shape_patch_offset
batch_patches_remaining = pred.size(0)-batch_offset
# append estimated patch properties batch to properties for the current shape
shape_properties[shape_patch_offset:shape_patch_offset+min(shape_patches_remaining, batch_patches_remaining), :] = pred[
batch_offset:batch_offset+min(shape_patches_remaining, batch_patches_remaining), :]
batch_offset = batch_offset + min(shape_patches_remaining, batch_patches_remaining)
shape_patch_offset = shape_patch_offset + min(shape_patches_remaining, batch_patches_remaining)
if shape_patches_remaining <= batch_patches_remaining:
# save shape properties to disk
prop_saved = [False]*len(trainopt.outputs)
# save normals
oi = [i for i, o in enumerate(trainopt.outputs) if o in ['unoriented_normals', 'oriented_normals']]
if len(oi) > 1:
raise ValueError('Duplicate normal output.')
elif len(oi) == 1:
oi = oi[0]
normal_prop = shape_properties[:, output_pred_ind[oi]:output_pred_ind[oi]+3]
np.savetxt(os.path.join(model_outdir, dataset.shape_names[shape_ind]+'.normals'), normal_prop.cpu().numpy())
prop_saved[oi] = True
# save curvatures
oi1 = [i for i, o in enumerate(trainopt.outputs) if o == 'max_curvature']
oi2 = [i for i, o in enumerate(trainopt.outputs) if o == 'min_curvature']
if len(oi1) > 1 or len(oi2) > 1:
raise ValueError('Duplicate minimum or maximum curvature output.')
elif len(oi1) == 1 or len(oi2) == 1:
curv_prop = shape_properties.new_zeros(shape_properties.size(0), 2)
if len(oi1) == 1:
oi1 = oi1[0]
curv_prop[:, 0] = shape_properties[:, output_pred_ind[oi1]]
prop_saved[oi1] = True
if len(oi2) == 1:
oi2 = oi2[0]
curv_prop[:, 1] = shape_properties[:, output_pred_ind[oi2]]
prop_saved[oi2] = True
np.savetxt(os.path.join(model_outdir, dataset.shape_names[shape_ind]+'.curv'), curv_prop.cpu().numpy())
if not all(prop_saved):
raise ValueError('Not all shape properties were saved, some of them seem to be unsupported.')
# save point indices
if opt.sampling != 'full':
np.savetxt(os.path.join(model_outdir, dataset.shape_names[shape_ind]+'.idx'), datasampler.shape_patch_inds[shape_ind], fmt='%d')
# start new shape
if shape_ind + 1 < len(dataset.shape_names):
shape_patch_offset = 0
shape_ind = shape_ind + 1
if opt.sampling == 'full':
shape_patch_count = dataset.shape_patch_count[shape_ind]
elif opt.sampling == 'sequential_shapes_random_patches':
# shape_patch_count = min(opt.patches_per_shape, dataset.shape_patch_count[shape_ind])
shape_patch_count = len(datasampler.shape_patch_inds[shape_ind])
else:
raise ValueError('Unknown sampling strategy: %s' % opt.sampling)
shape_properties = shape_properties.new_zeros(shape_patch_count, pred_dim)
def train_pcpnet(opt):
device = torch.device("cpu" if opt.gpu_idx < 0 else "cuda:%d" %
opt.gpu_idx)
# colored console output
green = lambda x: '\033[92m' + x + '\033[0m'
blue = lambda x: '\033[94m' + x + '\033[0m'
log_dirname = os.path.join(opt.logdir, opt.name)
params_filename = os.path.join(opt.outdir, '%s_params.pth' % (opt.name))
model_filename = os.path.join(opt.outdir, '%s_model.pth' % (opt.name))
desc_filename = os.path.join(opt.outdir, '%s_description.txt' % (opt.name))
if os.path.exists(log_dirname) or os.path.exists(model_filename):
response = input(
'A training run named "%s" already exists, overwrite? (y/n) ' %
(opt.name))
if response == 'y':
if os.path.exists(log_dirname):
shutil.rmtree(os.path.join(opt.logdir, opt.name))
else:
sys.exit()
# get indices in targets and predictions corresponding to each output
target_features = []
output_target_ind = []
output_pred_ind = []
output_loss_weight = []
pred_dim = 0
for o in opt.outputs:
if o == 'unoriented_normals' or o == 'oriented_normals':
if 'normal' not in target_features:
target_features.append('normal')
output_target_ind.append(target_features.index('normal'))
output_pred_ind.append(pred_dim)
output_loss_weight.append(1.0)
pred_dim += 3
elif o == 'max_curvature' or o == 'min_curvature':
if o not in target_features:
target_features.append(o)
output_target_ind.append(target_features.index(o))
output_pred_ind.append(pred_dim)
if o == 'max_curvature':
output_loss_weight.append(0.7)
else:
output_loss_weight.append(0.3)
pred_dim += 1
else:
raise ValueError('Unknown output: %s' % (o))
if pred_dim <= 0:
raise ValueError('Prediction is empty for the given outputs.')
# create model
if len(opt.patch_radius) == 1:
pcpnet = PCPNet(num_points=opt.points_per_patch,
output_dim=pred_dim,
use_point_stn=opt.use_point_stn,
use_feat_stn=opt.use_feat_stn,
sym_op=opt.sym_op,
point_tuple=opt.point_tuple)
else:
pcpnet = MSPCPNet(num_scales=len(opt.patch_radius),
num_points=opt.points_per_patch,
output_dim=pred_dim,
use_point_stn=opt.use_point_stn,
use_feat_stn=opt.use_feat_stn,
sym_op=opt.sym_op,
point_tuple=opt.point_tuple)
if opt.refine != '':
pcpnet.load_state_dict(torch.load(opt.refine))
if opt.seed < 0:
opt.seed = random.randint(1, 10000)
print("Random Seed: %d" % (opt.seed))
random.seed(opt.seed)
torch.manual_seed(opt.seed)
# create train and test dataset loaders
train_dataset = PointcloudPatchDataset(
root=opt.indir,
shape_list_filename=opt.trainset,
patch_radius=opt.patch_radius,
points_per_patch=opt.points_per_patch,
patch_features=target_features,
point_count_std=opt.patch_point_count_std,
seed=opt.seed,
identical_epochs=opt.identical_epochs,
use_pca=opt.use_pca,
center=opt.patch_center,
point_tuple=opt.point_tuple,
cache_capacity=opt.cache_capacity)
if opt.training_order == 'random':
train_datasampler = RandomPointcloudPatchSampler(
train_dataset,
patches_per_shape=opt.patches_per_shape,
seed=opt.seed,
identical_epochs=opt.identical_epochs)
elif opt.training_order == 'random_shape_consecutive':
train_datasampler = SequentialShapeRandomPointcloudPatchSampler(
train_dataset,
patches_per_shape=opt.patches_per_shape,
seed=opt.seed,
identical_epochs=opt.identical_epochs)
else:
raise ValueError('Unknown training order: %s' % (opt.training_order))
train_dataloader = torch.utils.data.DataLoader(train_dataset,
sampler=train_datasampler,
batch_size=opt.batchSize,
num_workers=int(
opt.workers))
test_dataset = PointcloudPatchDataset(
root=opt.indir,
shape_list_filename=opt.testset,
patch_radius=opt.patch_radius,
points_per_patch=opt.points_per_patch,
patch_features=target_features,
point_count_std=opt.patch_point_count_std,
seed=opt.seed,
identical_epochs=opt.identical_epochs,
use_pca=opt.use_pca,
center=opt.patch_center,
point_tuple=opt.point_tuple,
cache_capacity=opt.cache_capacity)
if opt.training_order == 'random':
test_datasampler = RandomPointcloudPatchSampler(
test_dataset,
patches_per_shape=opt.patches_per_shape,
seed=opt.seed,
identical_epochs=opt.identical_epochs)
elif opt.training_order == 'random_shape_consecutive':
test_datasampler = SequentialShapeRandomPointcloudPatchSampler(
test_dataset,
patches_per_shape=opt.patches_per_shape,
seed=opt.seed,
identical_epochs=opt.identical_epochs)
else:
raise ValueError('Unknown training order: %s' % (opt.training_order))
test_dataloader = torch.utils.data.DataLoader(test_dataset,
sampler=test_datasampler,
batch_size=opt.batchSize,
num_workers=int(opt.workers))
# keep the exact training shape names for later reference
opt.train_shapes = train_dataset.shape_names
opt.test_shapes = test_dataset.shape_names
print(
'training set: %d patches (in %d batches) - test set: %d patches (in %d batches)'
% (len(train_datasampler), len(train_dataloader),
len(test_datasampler), len(test_dataloader)))
try:
os.makedirs(opt.outdir)
except OSError:
pass
train_writer = SummaryWriter(os.path.join(log_dirname, 'train'))
test_writer = SummaryWriter(os.path.join(log_dirname, 'test'))
optimizer = optim.SGD(pcpnet.parameters(),
lr=opt.lr,
momentum=opt.momentum)
scheduler = lr_scheduler.MultiStepLR(
optimizer, milestones=[],
gamma=0.1) # milestones in number of optimizer iterations
pcpnet.to(device)
train_num_batch = len(train_dataloader)
test_num_batch = len(test_dataloader)
# save parameters
torch.save(opt, params_filename)
# save description
with open(desc_filename, 'w+') as text_file:
print(opt.desc, file=text_file)
for epoch in range(opt.nepoch):
train_batchind = -1
train_fraction_done = 0.0
train_enum = enumerate(train_dataloader, 0)
test_batchind = -1
test_fraction_done = 0.0
test_enum = enumerate(test_dataloader, 0)
for train_batchind, data in train_enum:
# update learning rate
scheduler.step(epoch * train_num_batch + train_batchind)
# set to training mode
pcpnet.train()
# get trainingset batch and upload to GPU
points = data[0]
target = data[1:-1]
points = points.transpose(2, 1)
points = points.to(device)
target = tuple(t.to(device) for t in target)
# zero gradients
optimizer.zero_grad()
# forward pass
pred, trans, _, _ = pcpnet(points)
loss = compute_loss(pred=pred,
target=target,
outputs=opt.outputs,
output_pred_ind=output_pred_ind,
output_target_ind=output_target_ind,
output_loss_weight=output_loss_weight,
patch_rot=trans if opt.use_point_stn else None,
normal_loss=opt.normal_loss)
# backpropagate through entire network to compute gradients of loss w.r.t. parameters
loss.backward()
# parameter optimization step
optimizer.step()
train_fraction_done = (train_batchind + 1) / train_num_batch
# print info and update log file
print('[%s %d: %d/%d] %s loss: %f' %
(opt.name, epoch, train_batchind, train_num_batch - 1,
green('train'), loss.item()))
train_writer.add_scalar('loss', loss.item(),
(epoch + train_fraction_done) *
train_num_batch * opt.batchSize)
while test_fraction_done <= train_fraction_done and test_batchind + 1 < test_num_batch:
# set to evaluation mode
pcpnet.eval()
test_batchind, data = next(test_enum)
# get testset batch and upload to GPU
points = data[0]
target = data[1:-1]
points = points.transpose(2, 1)
points = points.to(device)
target = tuple(t.to(device) for t in target)
# forward pass
with torch.no_grad():
pred, trans, _, _ = pcpnet(points)
loss = compute_loss(
pred=pred,
target=target,
outputs=opt.outputs,
output_pred_ind=output_pred_ind,
output_target_ind=output_target_ind,
output_loss_weight=output_loss_weight,
patch_rot=trans if opt.use_point_stn else None,
normal_loss=opt.normal_loss)
test_fraction_done = (test_batchind + 1) / test_num_batch
# print info and update log file
print('[%s %d: %d/%d] %s loss: %f' %
(opt.name, epoch, train_batchind, train_num_batch - 1,
blue('test'), loss.item()))
test_writer.add_scalar('loss', loss.item(),
(epoch + test_fraction_done) *
train_num_batch * opt.batchSize)
# save model, overwriting the old model
if epoch % opt.saveinterval == 0 or epoch == opt.nepoch - 1:
torch.save(pcpnet.state_dict(), model_filename)
# save model in a separate file in epochs 0,5,10,50,100,500,1000, ...
if epoch % (5 * 10**math.floor(math.log10(max(2, epoch - 1)))
) == 0 or epoch % 100 == 0 or epoch == opt.nepoch - 1:
torch.save(
pcpnet.state_dict(),
os.path.join(opt.outdir,
'%s_model_%d.pth' % (opt.name, epoch)))