def train(train_dataloader, test_dataloader, model, criterion, optimizer, lr_scheduler, bnm_scheduler, args, num_batch): PointcloudScaleAndTranslate = d_utils.PointcloudScaleAndTranslate( ) # initialize augmentation global g_acc g_acc = 0.91 # only save the model whose acc > 0.91 batch_count = 0 model.train() for epoch in range(args.epochs): for i, data in enumerate(train_dataloader, 0): if lr_scheduler is not None: lr_scheduler.step(epoch) if bnm_scheduler is not None: bnm_scheduler.step(epoch - 1) points, target = data points, target = points.cuda(), target.cuda() points, target = Variable(points), Variable(target) # farthest point sampling # fps_idx = pointnet2_utils.furthest_point_sample(points, 1200) # (B, npoint) # random sampling fps_idx = np.random.randint(0, points.shape[1] - 1, size=[points.shape[0], 1200]) fps_idx = torch.from_numpy(fps_idx).type(torch.IntTensor).cuda() fps_idx = fps_idx[:, np.random.choice(1200, args.num_points, False)] points = pointnet2_utils.gather_operation( points.transpose(1, 2).contiguous(), fps_idx).transpose(1, 2).contiguous() # (B, N, 3) # augmentation points.data = PointcloudScaleAndTranslate(points.data) optimizer.zero_grad() pred = model(points) target = target.view(-1) loss = criterion(pred, target) loss.backward() optimizer.step() if i % args.print_freq_iter == 0: print( '[epoch %3d: %3d/%3d] \t train loss: %0.6f \t lr: %0.5f' % (epoch + 1, i, num_batch, loss.data.clone(), lr_scheduler.get_lr()[0])) batch_count += 1 # validation in between an epoch if args.evaluate and batch_count % int( args.val_freq_epoch * num_batch) == 0: validate(test_dataloader, model, criterion, args, batch_count)
def train(train_dataloader, test_dataloader, model, criterion, optimizer, lr_scheduler, bnm_scheduler, args, num_batch): PointcloudScaleAndTranslate = d_utils.PointcloudScaleAndTranslate( ) # initialize augmentation global Class_mIoU, Inst_mIoU Class_mIoU, Inst_mIoU = 0.83, 0.85 batch_count = 0 #model.train() for epoch in range(args.epochs): model.train() losses = [] start_time = time.time() for i, data in enumerate(train_dataloader): points, target, cls = data target = target.cuda() points = points.cuda() one_hot_target = to_categorical(target, 50) print('target', target.shape) # augmentation points = PointcloudScaleAndTranslate(points) optimizer.zero_grad() batch_one_hot_cls = np.zeros((len(cls), 16)) # 16 object classes for b in range(len(cls)): batch_one_hot_cls[b, int(cls[b])] = 1 batch_one_hot_cls = torch.from_numpy( batch_one_hot_cls).float().cuda() # batch_one_hot_cls = Variable(batch_one_hot_cls.float().cuda()) pred, context = model(points, batch_one_hot_cls) pred = pred.view(-1, args.num_classes) target = target.view(-1, 1)[:, 0] loss = criterion(pred, target, None, context, one_hot_target) loss.backward() optimizer.step() losses.append(loss.item()) if lr_scheduler is not None: lr_scheduler.step(epoch) if bnm_scheduler is not None: bnm_scheduler.step(epoch - 1) #if i % args.print_freq_iter == 0: # print('[epoch %3d: %3d/%3d] \t train loss: %0.6f \t lr: %0.5f' %(epoch+1, i, num_batch, loss.data.clone(), lr_scheduler.get_lr()[0])) batch_count += 1 # validation in between an epoch #if (epoch < 3 or epoch > 40) and args.evaluate and batch_count % int(args.val_freq_epoch * num_batch) == 0: end_time = time.time() print('[epoch %3d time=%d s] \t train loss: %0.6f \t lr: %0.5f' % (epoch + 1, end_time - start_time, np.array(losses).mean(), lr_scheduler.get_lr()[0])) validate(test_dataloader, model, criterion, args, batch_count)
def train(train_dataloader, test_dataloader_z, test_dataloader_so3, model, criterion, optimizer, lr_scheduler, bnm_scheduler, args, num_batch): global Class_mIoU, Inst_mIoU Class_mIoU, Inst_mIoU = 0.75, 0.75 batch_count = 0 aug = d_utils.SO3Rotate() model.train() for epoch in range(args.epochs): for i, data in enumerate(train_dataloader, 0): if lr_scheduler is not None: lr_scheduler.step(epoch) if bnm_scheduler is not None: bnm_scheduler.step(epoch - 1) points, norm, target, cls = data points, norm, target = points.cuda(), norm.cuda(), target.cuda() # RS-CNN performs a translation to the input first if args.model == "rscnn_msn": points.data = d_utils.PointcloudScaleAndTranslate()( points.data) points.data, norm.data = aug(points.data, norm.data) optimizer.zero_grad() batch_one_hot_cls = np.zeros((len(cls), 16)) # 16 object classes for b in range(len(cls)): batch_one_hot_cls[b, int(cls[b])] = 1 batch_one_hot_cls = torch.from_numpy(batch_one_hot_cls) batch_one_hot_cls = batch_one_hot_cls.float().cuda() pred = model(points, norm, batch_one_hot_cls) pred = pred.view(-1, args.num_classes) target = target.view(-1, 1)[:, 0] loss = criterion(pred, target) loss.backward() optimizer.step() if i % args.print_freq_iter == 0: print( '[epoch %3d: %3d/%3d] \t train loss: %0.6f \t lr: %0.5f' % (epoch + 1, i, num_batch, loss.data.clone(), lr_scheduler.get_lr()[0])) batch_count += 1 if args.evaluate and batch_count % int( args.val_freq_epoch * num_batch) == 0: # validate(test_dataloader_z, model, criterion, args, batch_count, 'z') validate(test_dataloader_so3, model, criterion, args, batch_count, 'so3')
def train(train_dataloader, test_dataloader, model, criterion, optimizer, lr_scheduler, bnm_scheduler, args, num_batch): PointcloudScaleAndTranslate = d_utils.PointcloudScaleAndTranslate( ) # initialize augmentation global Class_mIoU, Inst_mIoU Class_mIoU, Inst_mIoU = 0.83, 0.85 batch_count = 0 model.train() for epoch in range(args.epochs): for i, data in enumerate(train_dataloader, 0): if lr_scheduler is not None: lr_scheduler.step(epoch) if bnm_scheduler is not None: bnm_scheduler.step(epoch - 1) points, target, cls = data points, target = points.cuda(), target.cuda() points, target = Variable(points), Variable(target) # augmentation points.data = PointcloudScaleAndTranslate(points.data) optimizer.zero_grad() batch_one_hot_cls = np.zeros((len(cls), 16)) # 16 object classes for b in range(len(cls)): batch_one_hot_cls[b, int(cls[b])] = 1 batch_one_hot_cls = torch.from_numpy(batch_one_hot_cls) batch_one_hot_cls = Variable(batch_one_hot_cls.float().cuda()) pred = model(points, batch_one_hot_cls) pred = pred.view(-1, args.num_classes) target = target.view(-1, 1)[:, 0] loss = criterion(pred, target) loss.backward() optimizer.step() if i % args.print_freq_iter == 0: print( '[epoch %3d: %3d/%3d] \t train loss: %0.6f \t lr: %0.5f' % (epoch + 1, i, num_batch, loss.data.clone(), lr_scheduler.get_lr()[0])) batch_count += 1 # validation in between an epoch if (epoch < 3 or epoch > 40) and args.evaluate and batch_count % int( args.val_freq_epoch * num_batch) == 0: validate(test_dataloader, model, criterion, args, batch_count)
def train(train_dataloader, test_dataloader_z, test_dataloader_so3, model, criterion, optimizer, lr_scheduler, bnm_scheduler, args, num_batch): aug = d_utils.ZRotate() global g_acc g_acc = 0.88 # only save the model whose acc > g_acc batch_count = 0 model.train() for epoch in range(args.epochs): for i, data in enumerate(train_dataloader, 0): if lr_scheduler is not None: lr_scheduler.step(epoch) if bnm_scheduler is not None: bnm_scheduler.step(epoch - 1) points, normals, target = data points, normals, target = points.cuda(), normals.cuda( ), target.cuda() if args.model == "pointnet2": fps_idx = pointnet2_utils.furthest_point_sample( points, 1024) # (B, npoint) points = pointnet2_utils.gather_operation( points.transpose(1, 2).contiguous(), fps_idx).transpose(1, 2).contiguous() # (B, N, 3) normals = pointnet2_utils.gather_operation( normals.transpose(1, 2).contiguous(), fps_idx).transpose(1, 2).contiguous() else: fps_idx = pointnet2_utils.furthest_point_sample( points, 1200) # (B, npoint) fps_idx = fps_idx[:, np.random.choice(1200, args.num_points, False )] points = pointnet2_utils.gather_operation( points.transpose(1, 2).contiguous(), fps_idx).transpose(1, 2).contiguous() # (B, N, 3) normals = pointnet2_utils.gather_operation( normals.transpose(1, 2).contiguous(), fps_idx).transpose(1, 2).contiguous() # # RS-CNN performs a translation to the input first points.data = d_utils.PointcloudScaleAndTranslate()( points.data) points.data, normals.data = aug(points.data, normals.data) optimizer.zero_grad() pred = model(points, normals) target = target.view(-1) loss = criterion(pred, target) loss.backward() optimizer.step() if i % args.print_freq_iter == 0: print( '[epoch %3d: %3d/%3d] \t train loss: %0.6f \t lr: %0.5f' % (epoch + 1, i, num_batch, loss.data.clone(), lr_scheduler.get_lr()[0])) batch_count += 1 # validation in between an epoch if args.evaluate and batch_count % int( args.val_freq_epoch * num_batch) == 0: # validate(test_dataloader_z, model, criterion, args, batch_count, 'z') validate(test_dataloader_so3, model, criterion, args, batch_count, 'so3')
def main(): args = parser.parse_args() with open(args.config) as f: config = yaml.load(f) print("\n**************************") for k, v in config['common'].items(): setattr(args, k, v) print('\n[%s]:' % (k), v) print("\n**************************\n") try: os.makedirs(args.save_path) except OSError: pass train_transforms = transforms.Compose([ d_utils.PointcloudToTensor(), d_utils.PointcloudScaleAndTranslate(), d_utils.PointcloudRandomInputDropout() ]) test_transforms = transforms.Compose([ d_utils.PointcloudToTensor(), #d_utils.PointcloudScaleAndTranslate() ]) train_dataset = ModelNet40Cls(num_points=args.num_points, root=args.data_root, transforms=train_transforms) train_dataloader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=int(args.workers)) test_dataset = ModelNet40Cls(num_points=args.num_points, root=args.data_root, transforms=test_transforms, train=False) test_dataloader = DataLoader(test_dataset, batch_size=args.batch_size, shuffle=False, num_workers=int(args.workers)) model = RSCNN_SSN(num_classes=args.num_classes, input_channels=args.input_channels, relation_prior=args.relation_prior, use_xyz=True) model.cuda() optimizer = optim.Adam(model.parameters(), lr=args.base_lr, weight_decay=args.weight_decay) lr_lbmd = lambda e: max(args.lr_decay**(e // args.decay_step), args.lr_clip / args.base_lr) bnm_lmbd = lambda e: max( args.bn_momentum * args.bn_decay** (e // args.decay_step), args.bnm_clip) lr_scheduler = lr_sched.LambdaLR(optimizer, lr_lbmd) bnm_scheduler = pt_utils.BNMomentumScheduler(model, bnm_lmbd) if args.checkpoint is not '': model.load_state_dict(torch.load(args.checkpoint)) print('Load model successfully: %s' % (args.checkpoint)) criterion = nn.CrossEntropyLoss() num_batch = len(train_dataset) / args.batch_size # training train(train_dataloader, test_dataloader, model, criterion, optimizer, lr_scheduler, bnm_scheduler, args, num_batch)
def train(ss_dataloader, train_dataloader, test_dataloader, encoder, decoer, optimizer, lr_scheduler, bnm_scheduler, args, num_batch, begin_epoch): PointcloudScaleAndTranslate = d_utils.PointcloudScaleAndTranslate( ) # initialize augmentation PointcloudRotate = d_utils.PointcloudRotate() metric_criterion = MetricLoss() chamfer_criterion = ChamferLoss() global svm_best_acc40 batch_count = 0 encoder.train() decoer.train() for epoch in range(begin_epoch, args.epochs): np.random.seed() for i, data in enumerate(ss_dataloader, 0): if lr_scheduler is not None: lr_scheduler.step(epoch) if bnm_scheduler is not None: bnm_scheduler.step(epoch - 1) points = data points = Variable(points.cuda()) # data augmentation sampled_points = 1200 has_normal = (points.size(2) > 3) if has_normal: normals = points[:, :, 3:6].contiguous() points = points[:, :, 0:3].contiguous() fps_idx = pointnet2_utils.furthest_point_sample( points, sampled_points) # (B, npoint) fps_idx = fps_idx[:, np.random.choice(sampled_points, args. num_points, False)] points_gt = pointnet2_utils.gather_operation( points.transpose(1, 2).contiguous(), fps_idx).transpose(1, 2).contiguous() # (B, N, 3) if has_normal: normals = pointnet2_utils.gather_operation( normals.transpose(1, 2).contiguous(), fps_idx) points = PointcloudScaleAndTranslate(points_gt.data) # optimize optimizer.zero_grad() features1, fuse_global, normals_pred = encoder(points) global_feature1 = features1[2].squeeze(2) refs1 = features1[0:2] recon1 = decoer(fuse_global).transpose(1, 2) # bs, np, 3 loss_metric = metric_criterion(global_feature1, refs1) loss_recon = chamfer_criterion(recon1, points_gt) if has_normal: loss_normals = NormalLoss(normals_pred, normals) else: loss_normals = normals_pred.new(1).fill_(0) loss = loss_recon + loss_metric + loss_normals loss.backward() optimizer.step() if i % args.print_freq_iter == 0: print( '[epoch %3d: %3d/%3d] \t metric/chamfer/normal loss: %0.6f/%0.6f/%0.6f \t lr: %0.5f' % (epoch + 1, i, num_batch, loss_metric.item(), loss_recon.item(), loss_normals.item(), lr_scheduler.get_lr()[0])) batch_count += 1 # validation if args.evaluate and batch_count % int( args.val_freq_epoch * num_batch) == 0: svm_acc40 = validate(train_dataloader, test_dataloader, encoder, args) save_dict = { 'epoch': epoch + 1, # after training one epoch, the start_epoch should be epoch+1 'optimizer_state_dict': optimizer.state_dict(), 'encoder_state_dict': encoder.state_dict(), 'decoder_state_dict': decoer.state_dict(), 'svm_best_acc40': svm_best_acc40, } checkpoint_name = './ckpts/' + args.name + '.pth' torch.save(save_dict, checkpoint_name) if svm_acc40 == svm_best_acc40: checkpoint_name = './ckpts/' + args.name + '_best.pth' torch.save(save_dict, checkpoint_name)
def validate(test_dataloader, model, criterion, args, iter): global Class_mIoU, Inst_mIoU, test_dataset model.eval() PointcloudScaleAndTranslate = d_utils.PointcloudScaleAndTranslate( ) # initialize augmentation seg_classes = test_dataset.seg_classes shape_ious = {cat: [] for cat in seg_classes.keys()} seg_label_to_cat = {} # {0:Airplane, 1:Airplane, ...49:Table} for cat in seg_classes.keys(): for label in seg_classes[cat]: seg_label_to_cat[label] = cat losses = 0.0 lens = len(test_dataloader) with torch.no_grad(): for _, data in enumerate(test_dataloader): points, target, cls = data target = target.cuda() points = points.cuda() one_hot_target = to_categorical(target, 50) # augmentation #points = PointcloudScaleAndTranslate(points) batch_one_hot_cls = np.zeros((len(cls), 16)) # 16 object classes for b in range(len(cls)): batch_one_hot_cls[b, int(cls[b])] = 1 batch_one_hot_cls = torch.from_numpy( batch_one_hot_cls).float().cuda() # batch_one_hot_cls = Variable(batch_one_hot_cls.float().cuda()) pred, context = model(points, batch_one_hot_cls) pred_t = pred.view(-1, args.num_classes) target_t = target.view(-1, 1)[:, 0] loss = criterion(pred_t, target_t, None, context, one_hot_target) losses += loss.item() """ points, target, cls = data #points, target = Variable(points, volatile=True), Variable(target, volatile=True) points, target = points.cuda(), target.cuda() one_hot_target=to_categorical(target,50) batch_one_hot_cls = np.zeros((len(cls), 16)) # 16 object classes for b in range(len(cls)): batch_one_hot_cls[b, int(cls[b])] = 1 batch_one_hot_cls = torch.from_numpy(batch_one_hot_cls).float().cuda() # batch_one_hot_cls = Variable(batch_one_hot_cls.float().cuda()) pred,context= model(points, batch_one_hot_cls) loss = criterion(pred.view(-1, args.num_classes), target.view(-1,1)[:,0],None,context,one_hot_target) losses+=loss.item() pred = pred.cpu() target = target.cpu() """ pred_val = torch.zeros(len(cls), args.num_points).type(torch.LongTensor) # pred to the groundtruth classes (selected by seg_classes[cat]) for b in range(len(cls)): cat = seg_label_to_cat[int(target[b, 0].cpu().numpy())] logits = pred[b, :, :] # (num_points, num_classes) pred_val[b, :] = logits[:, seg_classes[cat]].max( 1)[1] + seg_classes[cat][0] for b in range(len(cls)): segp = pred_val[b, :].cpu().numpy() segl = target[b, :].cpu().numpy() cat = seg_label_to_cat[int(segl[0])] part_ious = [0.0 for _ in range(len(seg_classes[cat]))] for l in seg_classes[cat]: if np.sum((segl == l) | (segp == l)) == 0: # part is not present in this shape part_ious[l - seg_classes[cat][0]] = 1.0 else: part_ious[l - seg_classes[cat][0]] = np.sum( (segl == l) & (segp == l)) / float( np.sum((segl == l) | (segp == l))) shape_ious[cat].append(np.mean(part_ious)) instance_ious = [] for cat in shape_ious.keys(): for iou in shape_ious[cat]: instance_ious.append(iou) shape_ious[cat] = np.mean(shape_ious[cat]) mean_class_ious = np.mean(list(shape_ious.values())) for cat in sorted(shape_ious.keys()): print('****** %s: %0.6f' % (cat, shape_ious[cat])) print('************ Test Loss: %0.6f' % (losses / lens)) print('************ Class_mIoU: %0.6f' % (mean_class_ious)) print('************ Instance_mIoU: %0.6f' % (np.mean(instance_ious))) if mean_class_ious > Class_mIoU or np.mean(instance_ious) > Inst_mIoU: if mean_class_ious > Class_mIoU: Class_mIoU = mean_class_ious if np.mean(instance_ious) > Inst_mIoU: Inst_mIoU = np.mean(instance_ious) torch.save( model.state_dict(), '%s/seg_msn_iter_%d_ins_%0.6f_cls_%0.6f.pth' % (args.save_path, iter, np.mean(instance_ious), mean_class_ious))