def prepare_model(option): loss = LineLoss(option.imagesize) dsac = DSAC(option.hypotheses, option.inlierthreshold, option.inlierbeta, option.inlieralpha, loss) point_nn = LineNN(option.capacity, option.receptivefield) if not option.cpu: point_nn = point_nn.cuda() point_nn.train() opt_point_nn = optim.Adam(point_nn.parameters(), lr=opt.learningrate) lrs_point_nn = optim.lr_scheduler.StepLR(opt_point_nn, opt.lrstep, gamma=0.5) return point_nn, dsac, opt_point_nn, lrs_point_nn
def eval_pcpnet(opt): if opt.distributed: torch.distributed.init_process_group(backend="nccl") # 2) 配置每个进程的gpu local_rank = torch.distributed.get_rank() torch.cuda.set_device(local_rank) device = torch.device("cuda", local_rank) print('stage 2 passed') else: device = torch.device("cpu" if opt.gpu_idx < 0 else "cuda:%d" % opt.gpu_idx) opt.models = opt.models.split() if opt.seed < 0: opt.seed = random.randint(1, 10000) 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) trainopt.indir2 = '../data2' if opt.distributed and local_rank == 0: print(trainopt) 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 = [] # trainopt.outputs = {'unoriented_normals'} for o in trainopt.outputs: if o == 'unoriented_normals' or o == 'oriented_normals': output_pred_ind.append(pred_dim) pred_dim += 3 else: raise ValueError('Unknown output: %s' % (o)) if len(trainopt.points_per_patch) == 1: if trainopt.generate_points_dim == 1: regressor = WDSAC( hyps=trainopt.hypotheses + 10, inlier_params=trainopt.inlier_params, patch_radius=trainopt.patch_radius, decoder=trainopt.decoder, use_mask=trainopt.use_mask, dim_pts=trainopt.in_points_dim, num_gpts=trainopt.generate_points_num, dim_gpts=trainopt.generate_points_dim, points_per_patch=trainopt.points_per_patch[0], sym_op=trainopt.sym_op, normal_loss=trainopt.normal_loss, seed=trainopt.seed, device=device, use_point_stn=trainopt.use_point_stn, use_feat_stn=trainopt.use_feat_stn) else: regressor = DSAC(hyps=trainopt.hypotheses + 10, inlier_params=trainopt.inlier_params, patch_radius=trainopt.patch_radius, decoder=trainopt.decoder, use_mask=trainopt.use_mask, dim_pts=trainopt.in_points_dim, num_gpts=trainopt.generate_points_num, dim_gpts=trainopt.generate_points_dim, points_per_patch=trainopt.points_per_patch[0], sym_op=trainopt.sym_op, normal_loss=trainopt.normal_loss, seed=trainopt.seed, device=device, use_point_stn=trainopt.use_point_stn, use_feat_stn=trainopt.use_feat_stn) else: regressor = MoEDSAC(hyps=trainopt.hypotheses + 10, inlier_params=trainopt.inlier_params, patch_radius=trainopt.patch_radius, share_pts_stn=trainopt.share_pts_stn, decoder=trainopt.decoder, use_mask=trainopt.use_mask, dim_pts=trainopt.in_points_dim, num_gpts=trainopt.generate_points_num, points_per_patch=trainopt.points_per_patch, sym_op=trainopt.sym_op, normal_loss=trainopt.normal_loss, seed=trainopt.seed, device=device, use_point_stn=trainopt.use_point_stn, use_feat_stn=trainopt.use_feat_stn) # if len(opt.expert_refine)>1: # dsac.refine(opt.expert_refine) regressor.to(device) if opt.distributed: if torch.cuda.device_count() > 1: print("Let's use", torch.cuda.device_count(), "GPUs!") regressor = torch.nn.parallel.DistributedDataParallel( regressor, device_ids=[local_rank], output_device=local_rank) #print(local_rank) state_dict = torch.load(model_filename, map_location=device) # print(state_dict.keys()) # print(regressor.state_dict().keys()) regressor.load_state_dict(state_dict) print('stage 3 passed') regressor.eval() dataset = PointcloudPatchDataset( root=opt.indir, root_in=trainopt.indir2, shape_list_filename=opt.dataset, patch_radius=trainopt.patch_radius, points_per_patch=trainopt.points_per_patch, dim_pts=trainopt.in_points_dim, knn=trainopt.knn, #patch_features=[], seed=opt.seed, 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) if opt.distributed: dataloader = torch.utils.data.DataLoader( dataset, sampler=DistributedSampler(datasampler), batch_size=model_batchSize, num_workers=int(opt.workers)) else: dataloader = torch.utils.data.DataLoader( dataset, sampler=datasampler, batch_size=model_batchSize, num_workers=int(opt.workers)) 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, 3, 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, target, _, dist = data points = points.transpose(2, 1) points = points.to(device) target = target.to(device) dist = dist.to(device) with torch.no_grad(): exp_loss, top_loss, pred, pts, _, _, _ = regressor( points, target, dist) 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()) print('saved normals for ' + dataset.shape_names[shape_ind]) prop_saved[oi] = True # save curvatures 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 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)) #print(trainopt.patch_radius) 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)) regressor = DSAC( trainopt.hypotheses, trainopt.inlierthreshold, trainopt.inlierbeta, trainopt.inlieralpha, trainopt.normal_loss, trainopt.seed,device, use_point_stn=trainopt.use_point_stn, use_feat_stn=trainopt.use_feat_stn, use_mask=trainopt.use_mask ) 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 #index =[] points, data_trans,mask_t = data points = points.transpose(2, 1) points = points.to(device) data_trans = data_trans.to(device) mask_t = mask_t.to(device) with torch.no_grad(): exp_loss, top_loss,normal,pts,mask = regressor(points,data_trans) viz = Visdom() assert viz.check_connection() Y=torch.zeros(512+32+1) Z=torch.zeros(512+1) Z[0:512]+=1 Z[512]+=2 Y[0:512]+=1 Y[512:512+32]+=2 Y[512+32]+=5 print(top_loss.mean()) for i in range(points.size(0)): #print("input",x[i].transpose(0,1)[0:100]) # print("predict",i,normal[i]) # print("ground truth",i,data_trans[i]) # print("top_loss_loss",i,top_loss[i],"\n") # print(mask_t[i]) # print(mask[i].view(-1)) viz.scatter( X=torch.cat((points[i].transpose(0,1),pts[i],torch.zeros(1,3).cuda()),0), Y=Y, opts=dict( title = str(i), #'legend': ['Men', 'Women'], markersize= 2, #markercolor=np.random.randint(0, 255, (3, 3,)), ) ) # # # viz.scatter( # # # X=torch.mul(points[i].transpose(0,1),mask[i]), # # # opts=dict( # # # title = str(i), # # # #'legend': ['Men', 'Women'], # # # markersize= 2, # # # #markercolor=np.random.randint(0, 255, (3, 3,)), # # # ) # # # ) # viz.scatter( # X=torch.cat((torch.mul(points[i].transpose(0,1),mask_t[i].view(-1,1)),torch.zeros(1,3).cuda(2)),0), # Y=Z, # opts=dict( # title = str(i)+"true", # #'legend': ['Men', 'Women'], # markersize= 2, # #markercolor=np.random.randint(0, 255, (3, 3,)), # ) # ) #print("pts",i,pts[i]) # # 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 < normal.size(0): shape_patches_remaining = shape_patch_count-shape_patch_offset batch_patches_remaining = normal.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), :] = normal[ 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(opt): # 1) 配置每个进程的gpu local_rank = torch.distributed.get_rank() torch.cuda.set_device(local_rank) device = torch.device("cuda", local_rank) print('stage 1 passed') # 2)初始化 # 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 local_rank == 0: if os.path.exists(log_dirname) or os.path.exists(model_filename): if os.path.exists(log_dirname): shutil.rmtree(os.path.join(opt.logdir, opt.name)) # if os.path.exists(log_dirname) or os.path.exists(model_filename): # if os.path.exists(log_dirname): # shutil.rmtree(os.path.join(opt.logdir, opt.name)) print('stage 2 passed') # 3) 封装之前要把模型移到对应的gpu if len(opt.points_per_patch) == 1: if opt.generate_points_dim == 1: dsac = WDSAC(hyps=opt.hypotheses, inlier_params=opt.inlier_params, patch_radius=opt.patch_radius, decoder=opt.decoder, use_mask=opt.use_mask, dim_pts=opt.in_points_dim, num_gpts=opt.generate_points_num, dim_gpts=opt.generate_points_dim, points_per_patch=opt.points_per_patch[0], sym_op=opt.sym_op, normal_loss=opt.normal_loss, seed=opt.seed, device=device, use_point_stn=opt.use_point_stn, use_feat_stn=opt.use_feat_stn) else: dsac = DSAC(hyps=opt.hypotheses, inlier_params=opt.inlier_params, patch_radius=opt.patch_radius, decoder=opt.decoder, use_mask=opt.use_mask, dim_pts=opt.in_points_dim, num_gpts=opt.generate_points_num, dim_gpts=opt.generate_points_dim, points_per_patch=opt.points_per_patch[0], sym_op=opt.sym_op, normal_loss=opt.normal_loss, seed=opt.seed, device=device, use_point_stn=opt.use_point_stn, use_feat_stn=opt.use_feat_stn) else: dsac = MoEDSAC(hyps=opt.hypotheses, inlier_params=opt.inlier_params, patch_radius=opt.patch_radius, share_pts_stn=opt.share_pts_stn, decoder=opt.decoder, use_mask=opt.use_mask, dim_pts=opt.in_points_dim, num_gpts=opt.generate_points_num, points_per_patch=opt.points_per_patch, sym_op=opt.sym_op, normal_loss=opt.normal_loss, seed=opt.seed, device=device, use_point_stn=opt.use_point_stn, use_feat_stn=opt.use_feat_stn) dsac.to(device) if torch.cuda.device_count() > 1: if opt.refine != '': # if the refine model is not from ddp dsac.load_state_dict(torch.load(opt.refine, map_location=device)) print("Let's use", torch.cuda.device_count(), "GPUs!") # 3) 封装 dsac = torch.nn.parallel.DistributedDataParallel( dsac, device_ids=[local_rank], output_device=local_rank) # if opt.refine != '': # if the refine model is not from ddp? # dsac.load_state_dict( torch.load(opt.refine, map_location=device) ) else: if opt.refine != '': dsac.load_state_dict(torch.load(opt.refine)) print('stage 3 passed') # 4)使用DistributedSampler train_dataset = PointcloudPatchDataset( root=opt.indir, root_in=opt.indir2, shape_list_filename=opt.trainset, patch_radius=opt.patch_radius, points_per_patch=opt.points_per_patch, dim_pts=opt.in_points_dim, knn=opt.knn, point_count_std=opt.patch_point_count_std, seed=opt.seed, identical_epochs=opt.identical_epochs, cache_capacity=opt.cache_capacity) train_datasampler = RandomPointcloudPatchSampler( train_dataset, patches_per_shape=opt.patches_per_shape, seed=opt.seed, identical_epochs=opt.identical_epochs) train_dataloader = DataLoader( train_dataset, sampler=DistributedSampler(train_datasampler), batch_size=opt.batchSize, num_workers=int(opt.workers)) test_dataset = PointcloudPatchDataset( root=opt.indir, root_in=opt.indir2, shape_list_filename=opt.testset, patch_radius=opt.patch_radius, points_per_patch=opt.points_per_patch, dim_pts=opt.in_points_dim, knn=opt.knn, seed=opt.seed, identical_epochs=opt.identical_epochs, cache_capacity=opt.cache_capacity) test_datasampler = RandomPointcloudPatchSampler( test_dataset, patches_per_shape=opt.patches_per_shape, seed=opt.seed, identical_epochs=opt.identical_epochs) test_dataloader = DataLoader(test_dataset, sampler=DistributedSampler(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 if local_rank == 0: train_writer = SummaryWriter(os.path.join(log_dirname, 'train')) test_writer = SummaryWriter(os.path.join(log_dirname, 'test')) lrate = opt.lr if opt.opti == 'SGD': optimizer = optim.SGD(dsac.parameters(), lr=lrate, momentum=opt.momentum) else: optimizer = optim.Adam(dsac.parameters(), lr=lrate) scheduler = lr_scheduler.MultiStepLR( optimizer, milestones=[1000], gamma=0.1) # milestones in number of optimizer iterations 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: # set to training mode dsac.train() points = data[0] #这时的point是64*512*3的类型 target = data[1] mask = data[2] dist = data[3] points = points.transpose(2, 1) points = points.to(device) target = target.to(device) mask = mask.to(device) dist = dist.to(device) # zero gradients #optimizer.zero_grad() exp_loss, top_loss, pred, pts, mask_p, patch_rot, _ = dsac( points, target, dist) loss = exp_loss.mean() # 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 tloss: %f loss: %f Top Loss:%f' % (opt.name, epoch, train_batchind, train_num_batch - 1, green('train'), loss, exp_loss.mean().item(), top_loss.mean())) if local_rank == 0: x1 = (epoch + train_fraction_done) * train_num_batch * opt.batchSize train_writer.add_scalars( 'loss', { 'meanLoss': exp_loss.mean().item(), 'topLoss': top_loss.mean().item() }, x1) while test_fraction_done <= train_fraction_done and test_batchind + 1 < test_num_batch: # set to evaluation mode dsac.eval() test_batchind, data = next(test_enum) points = data[0] #这时的point是64*512*3的类型 target = data[1] mask = data[2] dist = data[3] points = points.transpose(2, 1) points = points.to(device) target = target.to(device) mask = mask.to(device) dist = dist.to(device) # forward pass with torch.no_grad(): exp_loss, top_loss, pred, pts, mask_p, patch_rot, _ = dsac( points, target, dist) loss = exp_loss.mean() test_fraction_done = (test_batchind + 1) / test_num_batch # print info and update log file print('[%s %d: %d/%d] %s tloss: %f loss: %f Top Loss:%f' % (opt.name, epoch, train_batchind, train_num_batch - 1, blue('test'), loss, exp_loss.mean().item(), top_loss.mean())) if local_rank == 0: x1 = (epoch + test_fraction_done) * train_num_batch * opt.batchSize # test_writer.add_scalar('loss', exp_loss.mean().item(), x1) test_writer.add_scalars( 'loss', { 'meanLoss': exp_loss.mean().item(), 'topLoss': top_loss.mean().item() }, x1) # update learning rate scheduler.step() # save model, overwriting the old model if local_rank == 0: if epoch % opt.saveinterval == 0 or epoch == opt.nepoch - 1: print('model saved') torch.save(dsac.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( dsac.state_dict(), os.path.join(opt.outdir, '%s_model_%d.pth' % (opt.name, epoch)))
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)) #print(trainopt.patch_radius) dataset = PointcloudPatchDataset( root='/Users/jinwei/GItHub/DRNE/data/pclouds', root_in=trainopt.indir2, shape_list_filename=opt.dataset, patch_radius=trainopt.patch_radius, points_per_patch=trainopt.points_per_patch, dim_pts=trainopt.in_points_dim, knn=trainopt.knn, point_count_std=trainopt.patch_point_count_std, seed=trainopt.seed, identical_epochs=trainopt.identical_epochs, cache_capacity=trainopt.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)) regressor = DSAC(hyps=trainopt.hypotheses, inlier_params=trainopt.inlier_params, patch_radius=trainopt.patch_radius, decoder=trainopt.decoder, use_mask=trainopt.use_mask, dim_pts=trainopt.in_points_dim, num_gpts=trainopt.generate_points_num, dim_gpts=trainopt.generate_points_dim, points_per_patch=trainopt.points_per_patch[0], sym_op=trainopt.sym_op, normal_loss=trainopt.normal_loss, seed=trainopt.seed, device=device, use_point_stn=trainopt.use_point_stn, use_feat_stn=trainopt.use_feat_stn) regressor.load_state_dict( torch.load(model_filename, map_location=torch.device('cpu'))) 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[0] #这时的point是64*512*3的类型 target = data[1] mask = data[2] dist = data[3] points = points.transpose(2, 1) points = points.to(device) target = target.to(device) mask = mask.to(device) dist = dist.to(device) with torch.no_grad(): exp_loss, top_loss, pred, pts, mask_p, patch_rot, _ = regressor( points, target, dist) print(top_loss.mean()) pred_len = torch.max( torch.FloatTensor([sys.float_info.epsilon * 100]), pred.norm(p=2, dim=1, keepdim=True)) pred = pred / pred_len target_len = torch.max( torch.FloatTensor([sys.float_info.epsilon * 100]), target.norm(p=2, dim=2, keepdim=True)) target = target / target_len # plot a patch x, y = torch.meshgrid(torch.tensor([-10.0, 10.0]), torch.tensor([-10.0, 10.0])) for i in range(points.size(0)): pred_xy = pred[i, 0:2] / (pred[i, 2] + 1e-10) pred_xy = pred_xy.to(device) z = -(pred_xy[0] * x + pred_xy[1] * y) mlab.figure('patch_with_gpts', fgcolor=(0, 0, 0), bgcolor=(1, 1, 1)) mlab.points3d(10 * points[i, 0, :], 10 * points[i, 1, :], 10 * points[i, 2, :], color=(0.7, 0.7, 0.7), scale_factor=0.3, scale_mode='vector') mlab.points3d(10 * pts[i, :, 0], 10 * pts[i, :, 1], 10 * pts[i, :, 2], color=(0.2, 0.2, 0.2), scale_factor=0.7, scale_mode='vector') mlab.quiver3d(0.0, 0.0, 0.0, pred[i, 0], pred[i, 1], pred[i, 2], line_width=3, scale_factor=10, color=(0, 1, 0)) if (target[i, 0, :] - pred[i, :]).pow(2).sum() > ( target[i, 0, :] + pred[i, :]).pow(2).sum(): mlab.quiver3d(0.0, 0.0, 0.0, -target[i, 0, 0], -target[i, 0, 1], -target[i, 0, 2], line_width=3, scale_factor=10, color=(1, 0.0, 0.0)) else: mlab.quiver3d(0.0, 0.0, 0.0, target[i, 0, 0], target[i, 0, 1], target[i, 0, 2], line_width=3, scale_factor=10, color=(1, 0.0, 0.0)) mlab.surf(x, y, z, opacity=0.3) mlab.show()
parser.add_argument('--cpu', '-cpu', action='store_true', help='execute networks on CPU. Note that (RANSAC) line fitting anyway runs on CPU') parser.add_argument('--session', '-sid', default='', help='custom session name appended to output files. Useful to separate different runs of the program') opt = parser.parse_args() if len(opt.session) > 0: opt.session = '_' + opt.session sid = 'rf%d_c%d_h%d_t%.2f%s' % (opt.receptivefield, opt.capacity, opt.hypotheses, opt.inlierthreshold, opt.session) # setup the training process dataset = LineDataset(opt.imagesize, opt.imagesize) loss = LineLoss(opt.imagesize) dsac = DSAC(opt.hypotheses, opt.inlierthreshold, opt.inlierbeta, opt.inlieralpha, loss) # we train two CNNs in parallel # 1) a CNN that predicts points and is trained with DSAC -> PointNN (good idea) point_nn = LineNN(opt.capacity, opt.receptivefield) if not opt.cpu: point_nn = point_nn.cuda() point_nn.train() opt_point_nn = optim.Adam(point_nn.parameters(), lr=opt.learningrate) lrs_point_nn = optim.lr_scheduler.StepLR(opt_point_nn, opt.lrstep, gamma=0.5) # 2) a CNN that predicts the line parameters directly -> DirectNN (bad idea) direct_nn = LineNN(opt.capacity, 0, True) if not opt.cpu: direct_nn = direct_nn.cuda() direct_nn.train() opt_direct_nn = optim.Adam(direct_nn.parameters(), lr=opt.learningrate) lrs_direct_nn = optim.lr_scheduler.StepLR(opt_direct_nn, opt.lrstep, gamma=0.5)
help= 'custom session name appended to output files. Useful to separate different runs of the program' ) opt = parser.parse_args() if len(opt.session) > 0: opt.session = '_' + opt.session sid = 'rf%d_c%d_h%d_t%.2f%s' % (opt.receptivefield, opt.capacity, opt.hypotheses, opt.inlierthreshold, opt.session) # setup the training process dataset = CircleDataset(opt.imagesize, opt.imagesize) loss = CircleLoss(opt.imagesize) dsac = DSAC(opt.hypotheses, opt.inlierthreshold, opt.inlierbeta, opt.inlieralpha, loss) print("testing options", opt) print() print("testing soft inliner count") cX = 0.5 cY = 0.5 r = 0.5 x = torch.tensor([0.0, 0.5, 1.0, 0.96, 0.95, 0.94]) y = torch.tensor([0.5, 0.0, 0.5, 0.5, 0.5, 0.5]) result = dsac._DSAC__soft_inlier_count(cX, cY, r, x, y) print(result) print()
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 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)) regressor = DSAC( trainopt.hypotheses, trainopt.inlierthreshold, trainopt.inlierbeta, trainopt.inlieralpha, trainopt.normal_loss, trainopt.seed,device, use_point_stn=trainopt.use_point_stn, use_feat_stn=trainopt.use_feat_stn, use_mask=trainopt.use_mask ) 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, 3, 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(): exp_loss, top_loss,pred,pts,_ = regressor(points,data_trans) # 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] # # 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 # 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()) print('saved normals for ' + dataset.shape_names[shape_ind]) prop_saved[oi] = True # save curvatures 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_dsacpnet(opt): # gpu init # multi_gpus = False # if ',' in opt.gpu_idx: # gpu_ids = [int(id) for id in opt.gpu_idx.split(',')] # multi_gpus = True # else: # gpu_ids = [int(opt.gpu_idx)] # device = torch.device('cuda:{}'.format(gpu_ids[0]) if torch.cuda.is_available() else 'cpu') # 此处如果使用 下面一行代码,则会报错,RuntimeError: all tensors must be on devices[0] # 默认情况下 device 为0,因此需要指定 device id. device = torch.device("cpu" if opt.gpu_idx < 0 else "cuda:%d" % opt.gpu_idx) #if multi_gpus: # net = DataParallel(net, device_ids=gpu_ids).to(device) # margin = DataParallel(margin, device_ids=gpu_ids).to(device) #else: # net = net.to(device) # margin = margin.to(device) #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() criterion = nn.BCEWithLogitsLoss() dsac = DSAC(opt.hypotheses, opt.inlierthreshold, opt.inlierbeta, opt.inlieralpha, opt.normal_loss, opt.seed, device, use_point_stn=opt.use_point_stn, use_feat_stn=opt.use_feat_stn, use_mask=opt.use_mask, points_num=opt.points_num, points_per_patch=opt.points_per_patch, sym_op=opt.sym_op) 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, seed=opt.seed, identical_epochs=opt.identical_epochs, center=opt.patch_center, 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, seed=opt.seed, identical_epochs=opt.identical_epochs, center=opt.patch_center, 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(dsac.parameters(), lr=opt.lr, momentum=opt.momentum) scheduler = lr_scheduler.MultiStepLR( optimizer, milestones=[], gamma=0.1) # milestones in number of optimizer iterations #dsacpnet= torch.nn.DataParallel(dsacpnet, device_ids=gpu_ids).to(device) #dsac= torch.nn.DataParallel(dsac, device_ids=gpu_ids).to(device) dsac.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 dsac.train() points = data[0] #这时的point是64*512*3的类型 target = data[1] mask = data[2] points = points # for point in points: points = points.transpose(2, 1) points = points.to(device) target = target.to(device) mask = mask.to(device) # zero gradients optimizer.zero_grad() exp_loss, top_loss, _, pts, mask_p = dsac(points, target) if opt.use_mask: mask_p = mask_p.view(-1, opt.points_per_patch) mask_loss = criterion(mask_p, mask) # else: mask_loss = 0 exp_loss = exp_loss.mean() loss = exp_loss + mask_loss #+chamfer_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 tloss: %f loss: %f Top Loss:%f mask Loss:%f' % (opt.name, epoch, train_batchind, train_num_batch - 1, green('train'), loss, exp_loss.mean().item(), top_loss.mean(), mask_loss)) train_writer.add_scalar('loss', exp_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 dsac.eval() test_batchind, data = next(test_enum) points = data[0] #这时的point是64*512*3的类型 target = data[1] mask = data[2] points = points.transpose(2, 1) points = points.to(device) target = target.to(device) mask = mask.to(device) # forward pass with torch.no_grad(): exp_loss, top_loss, _, pts, mask_p = dsac(points, target) if opt.use_mask: mask_p = mask_p.view(-1, opt.points_per_patch) mask_loss = criterion(mask_p, mask) # else: mask_loss = 0 #mask_loss=criterion(mask_p, mask) loss = exp_loss + mask_loss test_fraction_done = (test_batchind + 1) / test_num_batch # print info and update log file print( '[%s %d: %d/%d] %s tloss: %f loss: %f Top Loss:%f mask Loss:%f ' % (opt.name, epoch, train_batchind, train_num_batch - 1, blue('test'), loss, exp_loss.mean().item(), top_loss.mean(), mask_loss)) test_writer.add_scalar('loss', exp_loss.mean().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(dsac.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( dsac.state_dict(), os.path.join(opt.outdir, '%s_model_%d.pth' % (opt.name, epoch)))