def main(args): def log_string(str): logger.info(str) print(str) '''HYPER PARAMETER''' os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu torch.backends.cudnn.enabled = False '''DATA LOADING''' TEST_DATASET = ModelNetDataLoader(root=args.data_root, tasks=args.test_tasks, labels=args.test_labels, partition='test', npoint=args.num_point, normal_channel=args.normal) testDataLoader = torch.utils.data.DataLoader(TEST_DATASET, batch_size=args.batch_size, shuffle=False, num_workers=2) '''MODEL LOADING''' num_class = 40 files = os.listdir(args.model_dir+'/logs') for f in files: if f.endswith('txt'): model_name = f.split('.')[0] MODEL = importlib.import_module(model_name) classifier = MODEL.get_model(num_class,normal_channel=args.normal).cuda() checkpoint = torch.load(str(args.model_dir) + '/checkpoints/best_model.pth') classifier.load_state_dict(checkpoint['model_state_dict']) with torch.no_grad(): instance_acc, class_acc = test(classifier.eval(), testDataLoader, vote_num=args.num_votes) print('Test Instance Accuracy: %f, Class Accuracy: %f' % (instance_acc, class_acc))
def evaluate(args): test_data, test_label = load_data( 'experiment/data/modelnet40_ply_hdf5_2048/', train=False) testDataset = ModelNetDataLoader(test_data, test_label) testDataLoader = torch.utils.data.DataLoader(testDataset, batch_size=args.batch_size, shuffle=False) log.debug('Building Model', args.model_name) if args.model_name == 'pointnet': num_class = 40 model = PointNetCls(num_class, args.feature_transform) else: model = PointNet2ClsMsg() torch.backends.cudnn.benchmark = True model = torch.nn.DataParallel(model).cuda() log.debug('Using gpu:', args.gpu) if args.pretrain is None: log.err('No pretrain model') return log.debug('Loading pretrain model...') state_dict = torch.load(args.pretrain) model.load_state_dict(state_dict) acc = test_clf(model.eval(), testDataLoader) log.msg(Test_Accuracy='%.5f' % (acc))
def run_trainer_evaluator(simclr, classifier, optimizer, args): data_train = ModelNetDataLoader(root=DATA_PATH, npoint=1024, split='train', normal_channel=False) data_test = ModelNetDataLoader(root=DATA_PATH, npoint=1024, split='test', normal_channel=False) best_vloss = INFINITE for epoch in range(args.epochs): print(f'Starting epoch [{epoch}/{args.epochs}]') # create the data loader for train and validation data train_loader = torch.utils.data.DataLoader(data_train, batch_size=args.batch_size, shuffle=True) test_loader = torch.utils.data.DataLoader(data_test, batch_size=args.batch_size, shuffle=False) # train and retrieve training loss t_loss, t_acc = train_validate(simclr, classifier, optimizer, train_loader, args, is_train=True) # retrieve validation loss v_loss, v_acc = train_validate(simclr, classifier, optimizer, test_loader, args, is_train=False) print( f'\nTotal epoch losses: train: {round(t_loss,4)} - validation: {round(v_loss,4)}' ) print( f'\nTotal epoch acc: train: {round(t_acc,4)} - validation: {round(v_acc,4)}\n', end='\r') # if the current loss is the new best, update checkpoint if v_loss < best_vloss: best_vloss = v_loss
def main(args): def log_string(str): logger.info(str) print(str) '''HYPER PARAMETER''' os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu '''CREATE DIR''' experiment_dir = 'log/classification/' + args.log_dir '''LOG''' args = parse_args() logger = logging.getLogger("Model") logger.setLevel(logging.INFO) formatter = logging.Formatter( '%(asctime)s - %(name)s - %(levelname)s - %(message)s') file_handler = logging.FileHandler('%s/eval.txt' % experiment_dir) file_handler.setLevel(logging.INFO) file_handler.setFormatter(formatter) logger.addHandler(file_handler) log_string('PARAMETER ...') log_string(args) '''DATA LOADING''' log_string('Load dataset ...') #分类使用的是ModelNet40数据集 data_path = '/home/wgk/dataset/Pointnet_Pointnet2_pytorch/modelnet40_normal_resampled/' test_dataset = ModelNetDataLoader(root=data_path, args=args, split='test', process_data=False) testDataLoader = torch.utils.data.DataLoader(test_dataset, batch_size=args.batch_size, shuffle=False, num_workers=10) '''MODEL LOADING''' num_class = args.num_category model_name = os.listdir(experiment_dir + '/logs')[0].split('.')[0] model = importlib.import_module(model_name) classifier = model.get_model(num_class, normal_channel=args.use_normals) if not args.use_cpu: classifier = classifier.cuda() checkpoint = torch.load( str(experiment_dir) + '/checkpoints/best_model.pth') classifier.load_state_dict(checkpoint['model_state_dict']) with torch.no_grad(): instance_acc, class_acc = test(classifier.eval(), testDataLoader, vote_num=args.num_votes, num_class=num_class) log_string('Test Instance Accuracy: %f, Class Accuracy: %f' % (instance_acc, class_acc))
def main(args): def log_string(str): logger.info(str) print(str) '''HYPER PARAMETER''' os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu '''CREATE DIR''' experiment_dir = 'log/classification/' + args.log_dir '''LOG''' args = parse_args() logger = logging.getLogger("Model") logger.setLevel(logging.INFO) formatter = logging.Formatter( '%(asctime)s - %(name)s - %(levelname)s - %(message)s') file_handler = logging.FileHandler('%s/eval.txt' % experiment_dir) file_handler.setLevel(logging.INFO) file_handler.setFormatter(formatter) logger.addHandler(file_handler) log_string('PARAMETER ...') log_string(args) '''DATA LOADING''' log_string('Load dataset ...') DATA_PATH = 'data/modelnet40_normal_resampled/' TEST_DATASET = ModelNetDataLoader(root=DATA_PATH, npoint=args.num_point, split='test', normal_channel=args.normal) testDataLoader = torch.utils.data.DataLoader(TEST_DATASET, batch_size=args.batch_size, shuffle=False, num_workers=4) '''MODEL LOADING''' num_class = 40 model_name = os.listdir(experiment_dir + '/logs')[0].split('.')[0] MODEL = importlib.import_module(model_name) classifier = MODEL.get_model(num_class, normal_channel=args.normal) checkpoint = torch.load(str(experiment_dir) + '/checkpoints/best_model.pth', map_location=torch.device('cpu')) classifier.load_state_dict(checkpoint['model_state_dict']) with torch.no_grad(): start_overall = time.time() instance_acc, class_acc = test(classifier.eval(), testDataLoader, vote_num=args.num_votes) log_string('Test Instance Accuracy: %f, Class Accuracy: %f' % (instance_acc, class_acc)) overall_time = time.time() - start_overall with open("overall.txt", 'a') as f: f.write(str(overall_time) + "\n")
def main(args): def log_string(str): logger.info(str) print(str) '''HYPER PARAMETER''' os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu '''CREATE DIR''' experiment_dir = args.log_dir '''LOG''' args = parse_args() logger = logging.getLogger("Model") logger.setLevel(logging.INFO) formatter = logging.Formatter( '%(asctime)s - %(name)s - %(levelname)s - %(message)s') file_handler = logging.FileHandler('%s/eval.txt' % experiment_dir) file_handler.setLevel(logging.INFO) file_handler.setFormatter(formatter) logger.addHandler(file_handler) log_string('PARAMETER ...') log_string(args) '''DATA LOADING''' log_string('Load dataset ...') DATA_PATH = 'data/modelnet40_normal_resampled/' TEST_DATASET = ModelNetDataLoader(root=DATA_PATH, npoint=args.num_point, split='test', normal_channel=args.normal) testDataLoader = torch.utils.data.DataLoader(TEST_DATASET, batch_size=args.batch_size, shuffle=False, num_workers=4) '''MODEL LOADING''' num_class = 40 model_name = "scratchpad__fixed_bug" MODEL = importlib.import_module(model_name) classifier = MODEL.get_model(num_class, normal_channel=args.normal).cuda() checkpoint = torch.load(str(experiment_dir) + '/seq_conv_best_model.pth') classifier.load_state_dict(checkpoint['model_state_dict']) with torch.no_grad(): instance_acc, class_acc = test(classifier.eval(), testDataLoader, vote_num=args.num_votes) log_string('Test Instance Accuracy: %f, Class Accuracy: %f' % (instance_acc, class_acc))
from models.pointnet import PointNetCls, feature_transform_regularizer from models.pointnet2 import PointNet2ClsMsg from models.dgcnn import DGCNN from models.pointcnn import PointCNNCls from utils import progress_bar, log_row sys.path.append("./emd/") import emd_module as emd device = torch.device("cuda" if torch.cuda.is_available() else "cpu") DATA_PATH = '/data/datasets/modelnet40_normal_resampled/' TEST_DATASET = ModelNetDataLoader(root=DATA_PATH, npoint=1024, split='test', normal_channel=False) test_loader = torch.utils.data.DataLoader(TEST_DATASET, batch_size=1, shuffle=False, num_workers=4, drop_last=False) PointcloudScaleAndTranslate = PointcloudScaleAndTranslate() print('======> Successfully loaded!') def cal_loss(pred, gold, smoothing=True): ''' Calculate cross entropy loss, apply label smoothing if needed. ''' gold = gold.contiguous().view(-1) if smoothing:
def main(args): config = yaml.load(open("./config/config.yaml", "r"), Loader=yaml.FullLoader) def log_string(str): logger.info(str) print(str) '''HYPER PARAMETER''' os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu '''CREATE DIR''' timestr = str(datetime.datetime.now().strftime('%Y-%m-%d_%H-%M')) experiment_dir = Path('./log/') experiment_dir.mkdir(exist_ok=True) experiment_dir = experiment_dir.joinpath('classification') experiment_dir.mkdir(exist_ok=True) if args.log_dir is None: experiment_dir = experiment_dir.joinpath(timestr) else: experiment_dir = experiment_dir.joinpath(args.log_dir) experiment_dir.mkdir(exist_ok=True) checkpoints_dir = experiment_dir.joinpath('checkpoints/') checkpoints_dir.mkdir(exist_ok=True) log_dir = experiment_dir.joinpath('logs/') log_dir.mkdir(exist_ok=True) '''LOG''' args = parse_args() logger = logging.getLogger("Model") logger.setLevel(logging.INFO) formatter = logging.Formatter( '%(asctime)s - %(name)s - %(levelname)s - %(message)s') file_handler = logging.FileHandler('%s/%s.txt' % (log_dir, args.model)) file_handler.setLevel(logging.INFO) file_handler.setFormatter(formatter) logger.addHandler(file_handler) log_string('PARAMETER ...') log_string(args) '''DATA LOADING''' log_string('Load dataset ...') DATA_PATH = 'data/modelnet40_normal_resampled/' TRAIN_DATASET = ModelNetDataLoader(root=DATA_PATH, npoint=args.num_point, split='train', normal_channel=args.normal) TEST_DATASET = ModelNetDataLoader(root=DATA_PATH, npoint=args.num_point, split='test', normal_channel=args.normal) trainDataLoader = torch.utils.data.DataLoader(TRAIN_DATASET, batch_size=args.batch_size, shuffle=True, num_workers=8) testDataLoader = torch.utils.data.DataLoader(TEST_DATASET, batch_size=args.batch_size, shuffle=False, num_workers=8) '''MODEL LOADING''' num_class = 40 MODEL = importlib.import_module(args.model) shutil.copy('./models/%s.py' % args.model, str(experiment_dir)) shutil.copy('./models/pointnet_util.py', str(experiment_dir)) # online network online_network = MODEL.get_model(num_class, normal_channel=args.normal).cuda() criterion = MODEL.get_loss().cuda() # predictor network device = 'cuda' if torch.cuda.is_available() else 'cpu' predictor = MLPHead( in_channels=online_network.projetion.net[-1].out_features, **config['network']['projection_head']).to(device) # target encoder target_network = MODEL.get_model(num_class, normal_channel=args.normal).cuda() # load pre-trained model if defined try: checkpoint = torch.load('checkpoints/model.pth') online_network.load_state_dict(checkpoint['online_network_state_dict']) target_network.load_state_dict(checkpoint['target_network_state_dict']) log_string('Use pretrain model') except: log_string('No existing model, starting training from scratch...') start_epoch = 0 # if args.optimizer == 'Adam': # optimizer = torch.optim.Adam( # classifier.parameters(), # lr=args.learning_rate, # betas=(0.9, 0.999), # eps=1e-08, # weight_decay=args.decay_rate # ) # else: # optimizer = torch.optim.SGD(classifier.parameters(), lr=0.01, momentum=0.9) optimizer = torch.optim.SGD( list(online_network.parameters()) + list(predictor.parameters()), **config['optimizer']['params']) trainer = BYOLTrainer(online_network=online_network, target_network=target_network, predictor=predictor, optimizer=optimizer, device=device, **config['trainer']) trainer.train_pointnet(trainDataLoader, testDataLoader)
def main(args): # 定义输出log和输出console def log_string(str): logger.info(str) print(str) '''HYPER PARAMETER''' os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu # 设置GPU的编号(可以多GPU运行) '''CREATE DIR''' timestr = str(datetime.datetime.now().strftime('%Y-%m-%d_%H-%M')) # 获取当前时间 experiment_dir = Path('./log/') # 对目录初始化path类 experiment_dir.mkdir(exist_ok=True) # 创建目录./log/,exist_ok=True目录存在不报错 # 创建classification目录 experiment_dir = experiment_dir.joinpath('classification') experiment_dir.mkdir(exist_ok=True) # 目录存在不报错 if args.log_dir is None: experiment_dir = experiment_dir.joinpath(timestr) else: experiment_dir = experiment_dir.joinpath(args.log_dir) experiment_dir.mkdir(exist_ok=True) checkpoints_dir = experiment_dir.joinpath( 'checkpoints/') # 创建checkpoints目录 checkpoints_dir.mkdir(exist_ok=True) log_dir = experiment_dir.joinpath('logs/') log_dir.mkdir(exist_ok=True) '''LOG''' args = parse_args() logger = logging.getLogger("Model") # 设置日志级别info : 打印info,warning,error,critical级别的日志 logger.setLevel(logging.INFO) # 配置日志的格式 formatter = logging.Formatter( '%(asctime)s - %(name)s - %(levelname)s - %(message)s') # 格式化字符串,路径为log_dir/args.model.txt file_handler = logging.FileHandler('%s/%s.txt' % (log_dir, args.model)) file_handler.setLevel(logging.INFO) file_handler.setFormatter(formatter) logger.addHandler(file_handler) log_string('PARAMETER ...') log_string(args) '''DATA LOADING''' log_string('Load dataset ...') DATA_PATH = 'data/modelnet40_normal_resampled/' TRAIN_DATASET = ModelNetDataLoader(root=DATA_PATH, npoint=args.num_point, split='train', normal_channel=args.normal) TEST_DATASET = ModelNetDataLoader(root=DATA_PATH, npoint=args.num_point, split='test', normal_channel=args.normal) import torch.utils.data.dataloader # 读取TRAIN_DATASET,设置batch_size,shuffle=True,打乱顺序,num_workers多线程 trainDataLoader = torch.utils.data.DataLoader(TRAIN_DATASET, batch_size=args.batch_size, shuffle=True, num_workers=4) # test数据集不需要打乱顺序 testDataLoader = torch.utils.data.DataLoader(TEST_DATASET, batch_size=args.batch_size, shuffle=False, num_workers=4) '''MODEL LOADING''' num_class = 40 # 导入model模型,相当于 import model MODEL = importlib.import_module(args.model) # 拷贝文件和权限,拷贝到experiment_dir='log\\classification\\pointnet2_cls_msg' shutil.copy('./models/%s.py' % args.model, str(experiment_dir)) shutil.copy('./models/pointnet_util.py', str(experiment_dir)) # 调用model(如pointnet_cls.py)中的方法 classifier = MODEL.get_model(num_class, normal_channel=args.normal).cuda() criterion = MODEL.get_loss().cuda() # 间断后继续训练 try: # 读取pth文件,pth中都以字典存储 # 将每一层与它的对应参数建立映射关系.(如model的每一层的weights及偏置等等)储存 # (注意,只有那些参数可以训练的layer才会被保存到模型的state_dict中,如卷积层,线性层等等) # 优化器对象Optimizer也有一个state_dict,它包含了优化器的状态以及被使用的超参数(如lr, momentum,weight_decay等) checkpoint = torch.load( str(experiment_dir) + '/checkpoints/best_model.pth') # 读取epoch start_epoch = checkpoint['epoch'] # 冲checkpoint中读取model_state_dict,并且用load_state_dict恢复模型参数 classifier.load_state_dict(checkpoint['model_state_dict']) # 写入log log_string('Use pretrain model') except: log_string('No existing model, starting training from scratch...') start_epoch = 0 # 创建optimizer优化器对象,这个对象能够保持当前参数状态并基于计算得到的梯度进行参数更新 if args.optimizer == 'Adam': optimizer = torch.optim.Adam( # 输入参数 classifier.parameters(), lr=args.learning_rate, betas=(0.9, 0.999), eps=1e-08, weight_decay=args.decay_rate) else: optimizer = torch.optim.SGD(classifier.parameters(), lr=0.01, momentum=0.9) # todo: https://blog.csdn.net/qyhaill/article/details/103043637 # 每过step_size个epoch,做一次更新 scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=20, gamma=0.7) global_epoch = 0 global_step = 0 best_instance_acc = 0.0 best_class_acc = 0.0 mean_correct = [] '''TRANING''' # 输出log logger.info('Start training...') for epoch in range(start_epoch, args.epoch): log_string('Epoch %d (%d/%s):' % (global_epoch + 1, epoch + 1, args.epoch)) # 更新lr,在epoch处调整 scheduler.step() # total迭代总次数,默认为迭代元素的长度,smoothing:0-1,0平均速度,1当前速度 for batch_id, data in tqdm(enumerate(trainDataLoader, 0), total=len(trainDataLoader), smoothing=0.9): # 从trainDataLoader中提取points:点集,target:对应的目标 points, target = data # 将tensor转化为numpy points = points.data.numpy() '''数据增强模块''' # 随机丢点 points = provider.random_point_dropout(points) # 随机范围 points[:, :, 0:3] = provider.random_scale_point_cloud(points[:, :, 0:3]) # 随机移动点云 points[:, :, 0:3] = provider.shift_point_cloud(points[:, :, 0:3]) # 转化为tensor points = torch.Tensor(points) # 得到target target = target[:, 0] # 输入要求(N,3)因此转置 points = points.transpose(2, 1) # 存入显存 points, target = points.cuda(), target.cuda() # 梯度归零 optimizer.zero_grad() # 训练模式 classifier = classifier.train() # 得到pointnet_cls前向传播返回的两个数据 pred, trans_feat = classifier(points) # loss loss = criterion(pred, target.long(), trans_feat) # 得到最大值的index pred_choice = pred.data.max(1)[1] # 将预测pred_choice和target(label)比较返回Boolean,将正确的求和算出预测正确的个数 correct = pred_choice.eq(target.long().data).cpu().sum() # correct.item()提取出tensor中的元素,将正确率其添加到mean_correct数组中 mean_correct.append(correct.item() / float(points.size()[0])) # 反向传播 loss.backward() # 更新梯度 optimizer.step() global_step += 1 train_instance_acc = np.mean(mean_correct) log_string('Train Instance Accuracy: %f' % train_instance_acc) # 在这个block(with torch.no_grad():)中不需要计算梯度(test模式) with torch.no_grad(): # classifier.eval():在test模式中禁用dropout和BN instance_acc, class_acc = test(classifier.eval(), testDataLoader) if (instance_acc >= best_instance_acc): best_instance_acc = instance_acc best_epoch = epoch + 1 if (class_acc >= best_class_acc): best_class_acc = class_acc log_string('Test Instance Accuracy: %f, Class Accuracy: %f' % (instance_acc, class_acc)) log_string('Best Instance Accuracy: %f, Class Accuracy: %f' % (best_instance_acc, best_class_acc)) if (instance_acc >= best_instance_acc): logger.info('Save model...') savepath = str(checkpoints_dir) + '/best_model.pth' log_string('Saving at %s' % savepath) state = { 'epoch': best_epoch, 'instance_acc': instance_acc, 'class_acc': class_acc, 'model_state_dict': classifier.state_dict(), 'optimizer_state_dict': optimizer.state_dict(), } torch.save(state, savepath) global_epoch += 1 logger.info('End of training...')
def main(args): '''HYPER PARAMETER''' os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu datapath = './data/ModelNet/' auxiliarypath = os.path.join(datapath, 'auxiliary') if ~os.path.exists(auxiliarypath): try: os.mkdir(auxiliarypath) except: pass '''CREATE DIR''' experiment_dir = Path('./experiment/') experiment_dir.mkdir(exist_ok=True) file_dir = Path( str(experiment_dir) + '/%sModelNet40-' % args.model_name + str(datetime.datetime.now().strftime('%Y-%m-%d_%H-%M'))) file_dir.mkdir(exist_ok=True) checkpoints_dir = file_dir.joinpath('checkpoints/') checkpoints_dir.mkdir(exist_ok=True) log_dir = file_dir.joinpath('logs/') log_dir.mkdir(exist_ok=True) '''LOG''' args = parse_args() logger = logging.getLogger(args.model_name) logger.setLevel(logging.INFO) formatter = logging.Formatter( '%(asctime)s - %(name)s - %(levelname)s - %(message)s') file_handler = logging.FileHandler( str(log_dir) + 'train_%s_cls.txt' % args.model_name) file_handler.setLevel(logging.INFO) file_handler.setFormatter(formatter) logger.addHandler(file_handler) logger.info( '---------------------------------------------------TRANING---------------------------------------------------' ) logger.info('PARAMETER ...') logger.info(args) '''MODEL LOADING''' num_class = 40 classifier = MainNet(num_class, normal=args.with_normal, random_sp=args.random_sample, random_nb=args.random_neighbor).cuda() if args.pretrain is not None: print('Use pretrain model...') logger.info('Use pretrain model') checkpoint = torch.load(args.pretrain) start_epoch = checkpoint['epoch'] classifier.load_state_dict(checkpoint['model_state_dict']) else: print('No existing model, starting training from scratch...') start_epoch = 0 if args.optimizer == 'SGD': optimizer = torch.optim.SGD(classifier.parameters(), lr=0.01, momentum=0.9) elif args.optimizer == 'Adam': optimizer = torch.optim.Adam(classifier.parameters(), lr=args.learning_rate, betas=(0.9, 0.999), eps=1e-08, weight_decay=args.decay_rate) scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=30, gamma=0.5) global_epoch = 0 global_step = 0 best_tst_accuracy = 0.0 blue = lambda x: '\033[94m' + x + '\033[0m' '''DATA LOADING''' logger.info('Load dataset ...') train_data, train_label, test_data, test_label = load_normal_data( './data/modelnet40_normal_resampled/') # logger.info('construct_MLS for train data...') # construct_MLS_multi(train_data, auxiliarypath, classifier.point_num) # logger.info('construct_MLS for test data...') # construct_MLS_multi(test_data, auxiliarypath, classifier.point_num, phase='test') train_neighbor_lists, train_data_idx_lists, train_local_axises = loadAuxiliaryInfo( auxiliarypath, classifier.point_num) logger.info("The number of training data is: %d", train_data.shape[0]) logger.info("The number of test data is: %d", test_data.shape[0]) trainDataset = ModelNetDataLoader(train_data, train_label, train_neighbor_lists, train_data_idx_lists, train_local_axises) test_neighbor_lists, test_data_idx_lists, test_local_axises = loadAuxiliaryInfo( auxiliarypath, classifier.point_num, phase='test') testDataset = ModelNetDataLoader(test_data, test_label, test_neighbor_lists, test_data_idx_lists, test_local_axises) trainDataLoader = torch.utils.data.DataLoader(trainDataset, batch_size=args.batchsize, shuffle=True) testDataLoader = torch.utils.data.DataLoader(testDataset, batch_size=args.batchsize, shuffle=False) '''TRANING''' logger.info('Start training...') first_time = True for epoch in range(start_epoch, args.epoch): print('Epoch %d (%d/%s):' % (global_epoch + 1, epoch + 1, args.epoch)) logger.info('Epoch %d (%d/%s):', global_epoch + 1, epoch + 1, args.epoch) scheduler.step() losses = 0 lc_stds = 0 lc_consistences = 0 mean_correct = [] for batch_id, data in tqdm(enumerate(trainDataLoader, 0), total=len(trainDataLoader), smoothing=0.9): points, target, neighbor_lists, data_idx_lists, local_axises = data # target = target[:, 0] # points = points.transpose(2, 1) if args.with_normal: points, target, neighbor_lists, data_idx_lists, local_axises = \ points.float().cuda(), target.cuda(), neighbor_lists.cuda(), data_idx_lists.cuda(), local_axises.cuda() else: points, target, neighbor_lists, data_idx_lists, local_axises = \ points[:,:,0:3].float().cuda(), target.cuda(), neighbor_lists.cuda(), data_idx_lists.cuda(), local_axises.cuda() optimizer.zero_grad() classifier = classifier.train() pred, lc_std, lc_consistence = classifier(points, neighbor_lists, data_idx_lists, local_axises) loss = F.nll_loss(pred, target.long()) total_loss = loss - lc_std - lc_consistence total_loss.backward() # print(classifier.axis_net.fc1.weight.grad) optimizer.step() global_step += 1 losses += loss.item() lc_stds += lc_std.item() lc_consistences += lc_consistence.item() pred_choice = pred.data.max(1)[1] correct = pred_choice.eq(target.long().data).cpu().sum() mean_correct.append(correct.item() / float(points.size()[0])) if epoch % 10 == 9: train_acc = test(classifier.eval(), trainDataLoader, args.with_normal) if args.train_metric else None writer3.add_scalar('quadratic', train_acc, global_step=epoch) acc = test(classifier.eval(), testDataLoader, args.with_normal) writer1.add_scalar('quadratic', losses / (batch_id + 1), global_step=epoch) writer2.add_scalar('quadratic', acc, global_step=epoch) print('\r Loss: %f' % float(losses / (batch_id + 1)), 'lc_std: %f' % float(lc_stds / (batch_id + 1)), 'lc_consistences: %f' % float(lc_consistences / (batch_id + 1)), 'acc: %f' % np.mean(mean_correct)) logger.info('Loss: %.2f', losses / (batch_id + 1)) if args.train_metric and epoch % 10 == 9: print('Train Accuracy: %f' % train_acc) logger.info('Train Accuracy: %f', (train_acc)) print( '\r Test %s: %f *** %s: %f' % (blue('Accuracy'), acc, blue('Best Accuracy'), best_tst_accuracy)) logger.info('Test Accuracy: %f *** Best Test Accuracy: %f', acc, best_tst_accuracy) if (acc >= best_tst_accuracy) and epoch > 10: best_tst_accuracy = acc logger.info('Save model...') save_checkpoint(global_epoch + 1, train_acc if args.train_metric else 0.0, acc, classifier, optimizer, str(checkpoints_dir), args.model_name) print('Saving model....') global_epoch += 1 print('Best Accuracy: %f' % best_tst_accuracy) logger.info('End of training...')
def main(args): '''HYPER PARAMETER''' os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu datapath = './data/ModelNet/' '''CREATE DIR''' experiment_dir = Path('./eval_experiment/') experiment_dir.mkdir(exist_ok=True) file_dir = Path( str(experiment_dir) + '/%sModelNet40-' % args.model_name + str(datetime.datetime.now().strftime('%Y-%m-%d_%H-%M'))) file_dir.mkdir(exist_ok=True) checkpoints_dir = file_dir.joinpath('checkpoints/') checkpoints_dir.mkdir(exist_ok=True) os.system('cp %s %s' % (args.kb1checkpoint, checkpoints_dir)) log_dir = file_dir.joinpath('logs/') log_dir.mkdir(exist_ok=True) '''LOG''' args = parse_args() logger = logging.getLogger(args.model_name) logger.setLevel(logging.INFO) formatter = logging.Formatter( '%(asctime)s - %(name)s - %(levelname)s - %(message)s') file_handler = logging.FileHandler( str(log_dir) + 'eval_%s_cls.txt' % args.model_name) file_handler.setLevel(logging.INFO) file_handler.setFormatter(formatter) logger.addHandler(file_handler) logger.info( '---------------------------------------------------EVAL---------------------------------------------------' ) logger.info('PARAMETER ...') logger.info(args) '''DATA LOADING''' logger.info('Load dataset ...') train_data, train_label, test_data, test_label = load_data( datapath, classification=True) logger.info("The number of training data is: %d", train_data.shape[0]) logger.info("The number of test data is: %d", test_data.shape[0]) testDataset = ModelNetDataLoader(test_data, test_label) testDataLoader = torch.utils.data.DataLoader(testDataset, batch_size=args.batchsize, shuffle=False) '''MODEL LOADING''' num_class = 39 kb1classifier = PointConvClsSsg(num_class).cuda() if args.kb1checkpoint is not None: print('Load k but 1 CheckPoint...') logger.info('Load k but 1 CheckPoint') kb1checkpoint = torch.load(args.kb1checkpoint) start_epoch = kb1checkpoint['epoch'] kb1classifier.load_state_dict(kb1checkpoint['model_state_dict']) else: print('Please load k but 1 Checkpoint to eval...') sys.exit(0) start_epoch = 0 num_class1 = 2 binaryclassifier = PointConvClsSsg(num_class1).cuda() if args.binarycheckpoint is not None: print('Load binary CheckPoint...') logger.info('Load binary CheckPoint') binarycheckpoint = torch.load(args.binarycheckpoint) start_epoch = binarycheckpoint['epoch'] binaryclassifier.load_state_dict(binarycheckpoint['model_state_dict']) else: print('Please load binary Checkpoint to eval...') sys.exit(0) start_epoch2 = 0 blue = lambda x: '\033[94m' + x + '\033[0m' '''EVAL''' logger.info('Start evaluating...') print('Start evaluating...') total_correct = 0 total_seen = 0 preds = [] for batch_id, data in tqdm(enumerate(testDataLoader, 0), total=len(testDataLoader), smoothing=0.9): pointcloud, target = data target = target[:, 0] #import ipdb; ipdb.set_trace() pred_view = torch.zeros(pointcloud.shape[0], num_class).cuda() binary_view = torch.zeros(pointcloud.shape[0], num_class1).cuda() for _ in range(args.num_view): pointcloud = generate_new_view(pointcloud) #import ipdb; ipdb.set_trace() #points = torch.from_numpy(pointcloud).permute(0, 2, 1) points = pointcloud.permute(0, 2, 1) points, target = points.cuda(), target.cuda() kb1classifier = kb1classifier.eval() binaryclassifier = binaryclassifier.eval() with torch.no_grad(): pred = kb1classifier(points) pred_binary = binaryclassifier(points) pred_view += pred binary_view += pred_binary kb1_logprob = pred_view.data binary_logprob = binary_view.data ## since we assigned the composite class the largest label, we will split the log-probability for the last label to two part, one for binary 0 and one for binary 1. binary_pred_logprob = kb1_logprob[:, -1].reshape( 1, len(kb1_logprob[:, -1])).transpose(0, 1).repeat(1, 2).view( -1, 2) + binary_logprob ## concatenate to get log-probability for all (40) classes pred_logprob = torch.from_numpy( np.c_[kb1_logprob[:, 0:-1].cpu().detach().numpy(), binary_pred_logprob.cpu().detach().numpy()]).to('cuda') pred_choices = pred_logprob.max(1)[1] ## reset labels mapper_dict = { **{key: key + 1 for key in range(12, 32)}, **{key: key + 2 for key in range(32, 38)}, **{ 38: 33, 39: 12 } } def mp(entry): return mapper_dict[entry] if entry in mapper_dict else entry mp = np.vectorize(mp) pred_choice = torch.from_numpy( np.array(mp(pred_choices.cpu().detach().numpy()))).to('cuda') preds.append(pred_choice.cpu().detach().numpy()) correct = pred_choice.eq( target.long().data).cpu().detach().numpy().sum() total_correct += correct.item() total_seen += float(points.size()[0]) accuracy = total_correct / total_seen ## confusion matrix cm = confusion_matrix(test_label.ravel(), np.concatenate(preds).ravel()) cm = cm.astype('float') / cm.sum(axis=1)[:, np.newaxis] t = pd.read_table('data/ModelNet/shape_names.txt', names=['label']) d = {key: val for key, val in zip(t.label, cm.diagonal())} print('Total Accuracy: %f' % accuracy) print('Accuracy per class:', d) logger.info('Total Accuracy: %f' % accuracy) logger.info('End of evaluation...')
def main(args): '''HYPER PARAMETER''' os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu datapath = './data/ModelNet/' '''CREATE DIR''' experiment_dir = Path('./experiment/') experiment_dir.mkdir(exist_ok=True) file_dir = Path( str(experiment_dir) + '/%sModelNet40-' % args.model_name + str(datetime.datetime.now().strftime('%Y-%m-%d_%H-%M'))) file_dir.mkdir(exist_ok=True) checkpoints_dir = file_dir.joinpath('checkpoints/') checkpoints_dir.mkdir(exist_ok=True) log_dir = file_dir.joinpath('logs/') log_dir.mkdir(exist_ok=True) '''LOG''' args = parse_args() logger = logging.getLogger(args.model_name) logger.setLevel(logging.INFO) formatter = logging.Formatter( '%(asctime)s - %(name)s - %(levelname)s - %(message)s') file_handler = logging.FileHandler( str(log_dir) + 'train_%s_cls.txt' % args.model_name) file_handler.setLevel(logging.INFO) file_handler.setFormatter(formatter) logger.addHandler(file_handler) logger.info( '---------------------------------------------------TRANING---------------------------------------------------' ) logger.info('PARAMETER ...') logger.info(args) '''DATA LOADING''' logger.info('Load dataset ...') train_data, train_label, test_data, test_label = load_data( datapath, classification=True) logger.info("The number of training data is: %d", train_data.shape[0]) logger.info("The number of test data is: %d", test_data.shape[0]) trainDataset = ModelNetDataLoader(train_data, train_label) testDataset = ModelNetDataLoader(test_data, test_label) trainDataLoader = torch.utils.data.DataLoader(trainDataset, batch_size=args.batchsize, shuffle=True) testDataLoader = torch.utils.data.DataLoader(testDataset, batch_size=args.batchsize, shuffle=False) '''MODEL LOADING''' num_class = 40 classifier = PointConvClsSsg(num_class).cuda() if args.pretrain is not None: print('Use pretrain model...') logger.info('Use pretrain model') checkpoint = torch.load(args.pretrain) start_epoch = checkpoint['epoch'] classifier.load_state_dict(checkpoint['model_state_dict']) else: print('No existing model, starting training from scratch...') start_epoch = 0 if args.optimizer == 'SGD': optimizer = torch.optim.SGD(classifier.parameters(), lr=0.01, momentum=0.9) elif args.optimizer == 'Adam': optimizer = torch.optim.Adam(classifier.parameters(), lr=args.learning_rate, betas=(0.9, 0.999), eps=1e-08, weight_decay=args.decay_rate) scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=30, gamma=0.5) global_epoch = 0 global_step = 0 best_tst_accuracy = 0.0 blue = lambda x: '\033[94m' + x + '\033[0m' '''TRANING''' logger.info('Start training...') for epoch in range(start_epoch, args.epoch): print('Epoch %d (%d/%s):' % (global_epoch + 1, epoch + 1, args.epoch)) logger.info('Epoch %d (%d/%s):', global_epoch + 1, epoch + 1, args.epoch) scheduler.step() for batch_id, data in tqdm(enumerate(trainDataLoader, 0), total=len(trainDataLoader), smoothing=0.9): points, target = data target = target[:, 0] points = points.transpose(2, 1) points, target = points.cuda(), target.cuda() optimizer.zero_grad() classifier = classifier.train() pred = classifier(points) loss = F.nll_loss(pred, target.long()) loss.backward() optimizer.step() global_step += 1 train_acc = test(classifier.eval(), trainDataLoader) if args.train_metric else None acc = test(classifier, testDataLoader) print('\r Loss: %f' % loss.data) logger.info('Loss: %.2f', loss.data) if args.train_metric: print('Train Accuracy: %f' % train_acc) logger.info('Train Accuracy: %f', (train_acc)) print( '\r Test %s: %f *** %s: %f' % (blue('Accuracy'), acc, blue('Best Accuracy'), best_tst_accuracy)) logger.info('Test Accuracy: %f *** Best Test Accuracy: %f', acc, best_tst_accuracy) if (acc >= best_tst_accuracy) and epoch > 5: best_tst_accuracy = acc logger.info('Save model...') save_checkpoint(global_epoch + 1, train_acc if args.train_metric else 0.0, acc, classifier, optimizer, str(checkpoints_dir), args.model_name) print('Saving model....') global_epoch += 1 print('Best Accuracy: %f' % best_tst_accuracy) logger.info('End of training...')
def main(args): def log_string(str): logger.info(str) print(str) '''HYPER PARAMETER''' os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu '''CREATE DIR''' timestr = str(datetime.datetime.now().strftime('%Y-%m-%d_%H-%M')) exp_dir = Path('./log/') exp_dir.mkdir(exist_ok=True) exp_dir = exp_dir.joinpath('classification') exp_dir.mkdir(exist_ok=True) if args.log_dir is None: exp_dir = exp_dir.joinpath(timestr) else: exp_dir = exp_dir.joinpath(args.log_dir) exp_dir.mkdir(exist_ok=True) checkpoints_dir = exp_dir.joinpath('checkpoints/') checkpoints_dir.mkdir(exist_ok=True) log_dir = exp_dir.joinpath('logs/') log_dir.mkdir(exist_ok=True) '''LOG''' args = parse_args() logger = logging.getLogger("Model") logger.setLevel(logging.INFO) formatter = logging.Formatter( '%(asctime)s - %(name)s - %(levelname)s - %(message)s') file_handler = logging.FileHandler('%s/%s.txt' % (log_dir, args.model)) file_handler.setLevel(logging.INFO) file_handler.setFormatter(formatter) logger.addHandler(file_handler) log_string('PARAMETER ...') log_string(args) '''DATA LOADING''' log_string('Load dataset ...') #设置数据集路径 data_path = '/home/wgk/dataset/Pointnet_Pointnet2_pytorch/modelnet40_normal_resampled/' #设置训练数据集 train_dataset = ModelNetDataLoader(root=data_path, args=args, split='train', process_data=args.process_data) #设置测试集 test_dataset = ModelNetDataLoader(root=data_path, args=args, split='test', process_data=args.process_data) #加载训练集合 trainDataLoader = torch.utils.data.DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=1, drop_last=True) testDataLoader = torch.utils.data.DataLoader(test_dataset, batch_size=args.batch_size, shuffle=False, num_workers=1) '''MODEL LOADING''' num_class = args.num_category #这里默认导入pointnet2_cls_msg.py文件,语法的使用 https://www.bilibili.com/read/cv5891176/ model = importlib.import_module(args.model) shutil.copy('./models/%s.py' % args.model, str(exp_dir)) shutil.copy('models/pointnet2_utils.py', str(exp_dir)) shutil.copy('./train_classification.py', str(exp_dir)) #加载分类器模型(实例化get_model这个class) classifier = model.get_model(num_class, normal_channel=args.use_normals) #实例化 criterion = model.get_loss() #.apply函数,是应用在 classifier.apply(inplace_relu) if not args.use_cpu: #将模型转移到gpu中 classifier = classifier.cuda() criterion = criterion.cuda() try: checkpoint = torch.load(str(exp_dir) + '/checkpoints/best_model.pth') start_epoch = checkpoint['epoch'] classifier.load_state_dict(checkpoint['model_state_dict']) log_string('Use pretrain model') except: log_string('No existing model, starting training from scratch...') start_epoch = 0 #配置优化器 if args.optimizer == 'Adam': optimizer = torch.optim.Adam(classifier.parameters(), lr=args.learning_rate, betas=(0.9, 0.999), eps=1e-08, weight_decay=args.decay_rate) else: optimizer = torch.optim.SGD(classifier.parameters(), lr=0.01, momentum=0.9) scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=20, gamma=0.7) global_epoch = 0 global_step = 0 best_instance_acc = 0.0 best_class_acc = 0.0 '''TRANING''' logger.info('Start training...') #默认训练200轮 for epoch in range(start_epoch, args.epoch): log_string('Epoch %d (%d/%s):' % (global_epoch + 1, epoch + 1, args.epoch)) mean_correct = [] classifier = classifier.train() scheduler.step() # for batch_id, (points, target) in tqdm(enumerate(trainDataLoader, 0), total=len(trainDataLoader), smoothing=0.9): optimizer.zero_grad() #获取到一个batch的数据,形状是(batch_size,点数=1024,channel) points = points.data.numpy() #print(points.shape) points = provider.random_point_dropout(points) #对数据集数据进行随机缩放 points[:, :, 0:3] = provider.random_scale_point_cloud(points[:, :, 0:3]) #对数据集数据进行随机旋转 points[:, :, 0:3] = provider.shift_point_cloud(points[:, :, 0:3]) #将ndarray数据转换为tensor points = torch.Tensor(points) #转置一下形状(batch_size,channel,点数) points = points.transpose(2, 1) if not args.use_cpu: #先将数据转移到显卡中 points, target = points.cuda(), target.cuda() #执行训练,返回预测的值 # pred.shape=[batchsize,40] trans_feat.shape=[batchsize,1024,1] pred, trans_feat = classifier(points) #print(pred.shape,trans_feat.shape) #loss函数按说应该是添加一个正则规范项,但是这里的loss实际上并没有添加 loss = criterion(pred, target.long(), trans_feat) pred_choice = pred.data.max(1)[1] #print(pred_choice) #查看预测值与真实值哪个位置相同,转到cpu中,并求总共有几个相同的 correct = pred_choice.eq(target.long().data).cpu().sum() #求正确预测的百分比,points.size()[0]=batchsize mean_correct.append(correct.item() / float(points.size()[0])) #这里的loss此时是一个tensor,所以可应用backward函数 loss.backward() optimizer.step() global_step += 1 train_instance_acc = np.mean(mean_correct) log_string('Train Instance Accuracy: %f' % train_instance_acc) with torch.no_grad(): instance_acc, class_acc = test(classifier.eval(), testDataLoader, num_class=num_class) if (instance_acc >= best_instance_acc): best_instance_acc = instance_acc best_epoch = epoch + 1 if (class_acc >= best_class_acc): best_class_acc = class_acc log_string('Test Instance Accuracy: %f, Class Accuracy: %f' % (instance_acc, class_acc)) log_string('Best Instance Accuracy: %f, Class Accuracy: %f' % (best_instance_acc, best_class_acc)) if (instance_acc >= best_instance_acc): logger.info('Save model...') savepath = str(checkpoints_dir) + '/best_model.pth' log_string('Saving at %s' % savepath) state = { 'epoch': best_epoch, 'instance_acc': instance_acc, 'class_acc': class_acc, 'model_state_dict': classifier.state_dict(), 'optimizer_state_dict': optimizer.state_dict(), } torch.save(state, savepath) global_epoch += 1 logger.info('End of training...')
def main(args): def log_string(str): logger.info(str) print(str) '''HYPER PARAMETER''' os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu '''SET THE SEED''' setup_seed(args.seed) '''CREATE DIR''' timestr = str(datetime.datetime.now().strftime('%Y-%m-%d_%H-%M')) experiment_dir = Path('./log/') experiment_dir.mkdir(exist_ok=True) experiment_dir = experiment_dir.joinpath('classification') experiment_dir.mkdir(exist_ok=True) if args.log_dir is None: experiment_dir = experiment_dir.joinpath(timestr) else: experiment_dir = experiment_dir.joinpath(args.log_dir) experiment_dir.mkdir(exist_ok=True) checkpoints_dir = experiment_dir.joinpath('checkpoints/') checkpoints_dir.mkdir(exist_ok=True) # log_dir = experiment_dir.joinpath('logs/') log_dir = experiment_dir.joinpath('./') log_dir.mkdir(exist_ok=True) '''LOG''' args = parse_args() logger = logging.getLogger("Model") logger.setLevel(logging.INFO) formatter = logging.Formatter( '%(asctime)s - %(name)s - %(levelname)s - %(message)s') file_handler = logging.FileHandler('%s/%s.txt' % (log_dir, args.model)) file_handler.setLevel(logging.INFO) file_handler.setFormatter(formatter) logger.addHandler(file_handler) log_string('PARAMETER ...') log_string(args) '''DATA TYPE''' if args.use_voxel: assert "voxel" in args.dataset assert "mink" in args.model assert args.voxel_size > 0 '''AUX SUPERVISION TYPE''' if args.aux == "pred": assert args.pred_path is not None if args.pred_path is not None: assert args.aux == "pred" args.with_pred = None args.with_instance = False args.with_seg = False if args.aux is not None: args.with_aux = True assert "scannet" in args.dataset if args.aux == "pred": args.with_pred = args.pred_path elif args.aux == "instance": args.with_instance = True elif args.aux == "seg": args.with_seg = True else: raise NotImplementedError else: args.with_aux = False '''DATA LOADING''' if "modelnet" in args.dataset: ''' the modelnet 40 loading, support both the point & ME-point Ver. ''' if not "voxel" in args.dataset: log_string('Load dataset {}'.format(args.dataset)) num_class = 40 DATA_PATH = './data/modelnet40_normal_resampled/' TRAIN_DATASET = ModelNetDataLoader(root=DATA_PATH, npoint=args.num_point, split='train', normal_channel=args.normal, apply_aug=True) TEST_DATASET = ModelNetDataLoader(root=DATA_PATH, npoint=args.num_point, split='test', normal_channel=args.normal) trainDataLoader = torch.utils.data.DataLoader( TRAIN_DATASET, batch_size=args.batch_size, shuffle=True, num_workers=args.num_worker) testDataLoader = torch.utils.data.DataLoader( TEST_DATASET, batch_size=args.batch_size, shuffle=False, num_workers=args.num_worker) else: assert args.dataset == 'modelnet_voxel' ''' use the modelnet example dataloader from the ME-engine however, it seems still is point-based, retrun point features, and feed in the TenesorField, not really the voxel-modelnet ''' log_string('Load dataset {}'.format(args.dataset)) num_class = 40 DATA_PATH = './data/modelnet40_ply_hdf5_2048' trainset = ModelNet40H5( phase="train", transform=CoordinateTransformation(trans=0.2), data_root=DATA_PATH, ) testset = ModelNet40H5( phase="test", transform=None, # no transform for test data_root=DATA_PATH, ) trainDataLoader = DataLoader( trainset, num_workers=args.num_worker, shuffle=True, batch_size=args.batch_size, collate_fn=minkowski_collate_fn, pin_memory=True, ) testDataLoader = DataLoader( testset, num_workers=args.num_worker, shuffle=False, batch_size=args.batch_size, collate_fn=minkowski_collate_fn, pin_memory=True, ) elif args.dataset == "scanobjnn": log_string('Load dataset {}'.format(args.dataset)) num_class = 15 DATA_PATH = './data/scanobjnn/main_split_nobg' TRAIN_DATASET = ScanObjectNNDataLoader(root=DATA_PATH, npoint=args.num_point, split='train', normal_channel=args.normal) TEST_DATASET = ScanObjectNNDataLoader(root=DATA_PATH, npoint=args.num_point, split='test', normal_channel=args.normal) trainDataLoader = torch.utils.data.DataLoader( TRAIN_DATASET, batch_size=args.batch_size, shuffle=True, num_workers=args.num_worker) testDataLoader = torch.utils.data.DataLoader( TEST_DATASET, batch_size=args.batch_size, shuffle=False, num_workers=args.num_worker) elif "scannet" in args.dataset: num_class = 21 if not "voxel" in args.dataset: if args.mode == "train": trainset = ScannetDataset( root='./data/scannet_v2/scannet_pickles', npoints=args.num_point, split='train', with_seg=args.with_seg, with_instance=args.with_instance, with_pred=args.pred_path) trainDataLoader = torch.utils.data.DataLoader( trainset, batch_size=args.batch_size, shuffle=True, num_workers=args.num_worker, pin_memory=True) if args.mode == 'export': final_trainset = ScannetDatasetWholeScene_evaluation(root='./data/scannet_v2/scannet_pickles', scene_list_dir='./data/scannet_v2/metadata',split='train',block_points=args.num_point, with_rgb=True, with_norm=True,\ with_seg=args.with_seg, with_instance=args.with_instance, with_pred=args.pred_path, delta=2.0) final_train_loader = torch.utils.data.DataLoader( final_trainset, batch_size=args.batch_size, shuffle=False, num_workers=0, pin_memory=True) final_testset = ScannetDatasetWholeScene_evaluation(root='./data/scannet_v2/scannet_pickles', scene_list_dir='./data/scannet_v2/metadata',split='eval',block_points=args.num_point, with_rgb=True, with_norm=True, \ with_seg=args.with_seg, with_instance=args.with_instance, with_pred=args.pred_path, delta=1.0) # DEBUG: change to 1.0 to axquire proper final_test_loader = torch.utils.data.DataLoader( final_testset, batch_size=args.batch_size, shuffle=False, num_workers=0, pin_memory=True) # generate the trainset as whole_dataset for export else: trainDataLoader = initialize_data_loader( DatasetClass=ScannetSparseVoxelizationDataset, data_root='data/scannet_v2/scannet_pickles', phase="train", threads=4, # num-workers shuffle=True, repeat=False, augment_data=True, batch_size=16, limit_numpoints=1200000, ) # TODO: the testloader else: raise NotImplementedError '''MODEL LOADING''' # copy files if args.mode == "train": if not os.path.exists(os.path.join(str(experiment_dir), 'model')): os.mkdir(os.path.join(str(experiment_dir), 'model')) for filename in os.listdir('./model'): if ".py" in filename: shutil.copy(os.path.join("./model", filename), os.path.join(str(experiment_dir), 'model')) shutil.copy("./train_cls.py", str(experiment_dir)) if "mink" not in args.model: # no use mink-net if "seg" in args.model: N = args.num_point else: N = args.num_point MODEL = importlib.import_module(args.model) classifier = MODEL.get_model(num_class, normal_channel=args.normal, N=N).cuda() criterion = MODEL.get_loss().cuda() classifier.loss = criterion else: ''' The Voxel-based Networks based on the MinkowskiEngine ''' # TODO: should align with above, using importlib.import_module, maybe fix later # classifier = ResNet14(in_channels=3, out_channels=num_class, D=3) # D is the conv spatial dimension, 3 menas 3-d shapes if "pointnet" in args.model: classifier = MinkowskiPointNet(in_channel=3, out_channel=41, embedding_channel=1024, dimension=3).cuda() elif "trans" in args.model: classifier = MinkowskiTransformer(in_channel=3, out_channel=41, num_class=num_class, embedding_channel=1024, dimension=3).cuda() criterion = nn.CrossEntropyLoss().cuda() classifier.loss = criterion '''Loading existing ckpt''' try: if args.pretrain: # FIXME: currently only loading the best_model.pth, should support string, maybe latter checkpoint = torch.load( str(experiment_dir) + '/checkpoints/best_model.pth') start_epoch = checkpoint['epoch'] classifier.load_state_dict(checkpoint['model_state_dict']) log_string('Use pretrain model') start_epoch = 0 else: start_epoch = 0 except: log_string('No existing model, starting training from scratch...') start_epoch = 0 if args.optimizer == 'Adam': optimizer = torch.optim.Adam(classifier.parameters(), lr=args.learning_rate, betas=(0.9, 0.999), eps=1e-08, weight_decay=args.decay_rate) elif args.optimizer == 'SGD': optimizer = torch.optim.SGD(classifier.parameters(), \ lr=args.learning_rate, momentum=0.9,\ weight_decay=args.decay_rate) else: raise NotImplementedError # scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=20, gamma=0.7) # Use MultiStepLR as in paper, decay by 10 at [120, 160] # scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=[120, 160], gamma=0.1) # FIXME: for scannet, now using the cosine anneal scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, args.epoch, eta_min=0.0) global_epoch = 0 global_step = 0 if "scannet" in args.dataset: best_mIoU = 0.0 else: best_instance_acc = 0.0 best_class_acc = 0.0 mean_correct = [] # only run for one epoch on the eval-only mode if args.mode == "eval" or args.mode == "export": assert args.pretrain start_epoch = 0 args.epoch = 1 # '''TRANING''' logger.info('Start training...') for epoch in range(start_epoch, args.epoch): log_string('Epoch %d (%d/%s):' % (global_epoch + 1, epoch + 1, args.epoch)) scheduler.step() log_string('Cur LR: {:.5f}'.format(optimizer.param_groups[0]['lr'])) # when eval only, skip the traininig part if args.mode == "train": '''The main training-loop''' for batch_id, data in tqdm(enumerate(trainDataLoader, 0), total=len(trainDataLoader), smoothing=0.9): if not args.use_voxel: if "modelnet" in args.dataset: # use points, normal unpacking points, target = data points = points.data.numpy() points = provider.random_point_dropout(points) # points[:,:, 0:3] = provider.random_scale_point_cloud(points[:,:, 0:3]) # points[:,:, 0:3] = provider.shift_point_cloud(points[:,:, 0:3]) points = torch.Tensor(points) target = target[:, 0] points = points.transpose(2, 1) points, target = points.cuda(), target.cuda() elif "scanobjnn" in args.dataset: points, target, mask = data points = points.data.numpy() # TODO: move the aug in the dataset but not here # points = provider.random_point_dropout(points) points[:, :, 0:3] = provider.random_scale_point_cloud( points[:, :, 0:3]) points[:, :, 0:3] = provider.shift_point_cloud(points[:, :, 0:3]) points = torch.Tensor(points) points = points.transpose(2, 1) points, target = points.cuda(), target.cuda() elif "scannet" in args.dataset: # TODO: fiil the scannet loading here # TODO: maybe implement the grad-accmu/or simply not if args.aux is not None: points, target, sample_weight, aux = data points, target, sample_weight, aux = points.float( ).transpose(1, 2).cuda(), target.cuda( ), sample_weight.cuda(), aux.cuda() else: points, target, sample_weight = data points, target, sample_weight = points.float( ).transpose( 1, 2).cuda(), target.cuda(), sample_weight.cuda() else: if "modelnet" in args.dataset: # use voxel # points = create_input_batch(data, True, 'cuda', quantization_size=args.voxel_size) data['coordinates'][:, 1:] = data[ 'coordinates'][:, 1:] / args.voxel_size points = ME.TensorField( coordinates=(data['coordinates'].cuda()), features=data['features'].cuda()) target = data['labels'].cuda() elif "scannet" in args.dataset: dat = ME.SparseTensor(features=data[1], coordinates=data[0]).cuda() target = data[2].cuda() optimizer.zero_grad() '''save the intermediate attention map''' # WANINIG: DISABLED FOR NOW!!! SAVE_INTERVAL = 50 NUM_PER_EPOCH = 1 if (epoch + 1) % SAVE_INTERVAL == 0: if batch_id < NUM_PER_EPOCH: classifier.save_flag = True elif batch_id == NUM_PER_EPOCH: intermediate_dict = classifier.save_intermediate() intermediate_path = os.path.join( experiment_dir, "attn") if not os.path.exists(intermediate_path): os.mkdir(intermediate_path) torch.save( intermediate_dict, os.path.join(intermediate_path, "epoch_{}".format(epoch))) log_string('Saved Intermediate at {}'.format(epoch)) else: classifier.save_flag = False else: classifier.save_flag = False classifier = classifier.train() # when with-instance, use instance label to guide the point-transformer training if args.aux is not None: pred = classifier(points, aux) else: pred = classifier(points) # if use_voxel, get the feature from the SparseTensor if args.use_voxel: pred = pred.F if 'scannet' in args.dataset: loss = criterion(pred, target.long(), sample_weight) else: loss = criterion(pred, target.long()) loss.backward() optimizer.step() global_step += 1 if "scannet" in args.dataset: pred_choice = torch.argmax(pred, dim=2).cpu().numpy() # B,N target = target.cpu().numpy() correct = np.sum(pred_choice == target) mean_correct.append(correct / pred_choice.size) else: pred_choice = pred.data.max(1)[1] correct = pred_choice.eq(target.long().data).cpu().sum() mean_correct.append(correct.item() / float(points.size()[0])) train_instance_acc = np.mean(mean_correct) log_string('Train Instance Accuracy: %f' % train_instance_acc) '''TEST''' if not "scannet" in args.dataset: # WARNING: Temporarily disable eval for scannet for now, just test at last if (epoch + 1) % 20 == 0: with torch.no_grad(): returned_metric = test(classifier.eval(), testDataLoader, num_class=num_class, log_string=log_string) if 'scannet' in args.dataset: mIoU = returned_metric if (mIoU >= best_mIoU): best_mIoU = mIoU best_epoch = epoch + 1 if (mIoU >= best_mIoU): logger.info('Save model...') savepath = str(checkpoints_dir) + '/best_model.pth' log_string('Saving at %s' % savepath) state = { 'epoch': best_epoch, 'mIoU': mIoU, 'model_state_dict': classifier.state_dict(), 'optimizer_state_dict': optimizer.state_dict(), } torch.save(state, savepath) else: instance_acc, class_acc = returned_metric if (instance_acc >= best_instance_acc): best_instance_acc = instance_acc best_epoch = epoch + 1 if (class_acc >= best_class_acc): best_class_acc = class_acc log_string( 'Test Instance Accuracy: %f, Class Accuracy: %f' % (instance_acc, class_acc)) log_string( 'Best Instance Accuracy: %f, Class Accuracy: %f' % (best_instance_acc, best_class_acc)) if (instance_acc >= best_instance_acc): logger.info('Save model...') savepath = str(checkpoints_dir) + '/best_model.pth' log_string('Saving at %s' % savepath) state = { 'epoch': best_epoch, 'instance_acc': instance_acc, 'class_acc': class_acc, 'model_state_dict': classifier.state_dict(), 'optimizer_state_dict': optimizer.state_dict(), } torch.save(state, savepath) global_epoch += 1 # final save of the model logger.info('Save model...') savepath = str(checkpoints_dir) + '/final_model.pth' log_string('Saving at %s' % savepath) state = { 'epoch': global_epoch, # 'mIoU': mIoU, 'model_state_dict': classifier.state_dict(), 'optimizer_state_dict': optimizer.state_dict(), } torch.save(state, savepath) # for the scannet dataset, test at last if args.dataset == 'scannet': if not os.path.exists(os.path.join(str(experiment_dir), 'pred')): os.mkdir(os.path.join(str(experiment_dir), 'pred')) if args.mode == "export": test_scannet(args, classifier.eval(), final_test_loader, log_string, with_aux=args.with_aux, save_dir=os.path.join(str(experiment_dir), 'pred'), split='eval') test_scannet(args, classifier.eval(), final_train_loader, log_string, with_aux=args.with_aux, save_dir=os.path.join(str(experiment_dir), 'pred'), split='train') else: test_scannet(args, classifier.eval(), final_test_loader, log_string, with_aux=args.with_aux, split='eval') # final save of the model logger.info('Save model...') savepath = str(checkpoints_dir) + '/best_model.pth' log_string('Saving at %s' % savepath) state = { 'epoch': global_epoch, # 'mIoU': mIoU, 'model_state_dict': classifier.state_dict(), 'optimizer_state_dict': optimizer.state_dict(), } torch.save(state, savepath) logger.info('End of training...')
def main(args): def log_string(str): logger.info(str) print(str) '''CUDA ENV SETTINGS''' if args.gpu is not None: os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu if args.cudnn_off: torch.backends.cudnn.enabled = False # needed on gypsum! # -------------------------------------------------------------------------- '''CREATE DIR''' # -------------------------------------------------------------------------- timestr = str(datetime.datetime.now().strftime('%Y-%m-%d_%H-%M')) experiment_dir = Path('./log/') experiment_dir.mkdir(exist_ok=True) experiment_dir = experiment_dir.joinpath('pretrain_part_seg') experiment_dir.mkdir(exist_ok=True) dir_name = args.model + '_ShapeNet' + \ '_k-%d_seed-%d_lr-%.6f_lr-step-%d_lr-decay-%.2f_wt-decay-%.6f_l2norm-%d' \ % ( args.k_shot, args.seed, args.learning_rate, args.step_size, args.lr_decay, args.decay_rate, int(args.l2_norm) ) if args.normal: dir_name = dir_name + '_normals' if args.selfsup: dir_name = dir_name + 'selfsup-%s_selfsup_margin-%.2f_lambda-%.2f' \ % (args.ss_dataset, args.margin, args.lmbda) if args.rotation_z: dir_name = dir_name + '_rotation-z' if args.rotation_z_45: dir_name = dir_name + '_rotation-z-45' if args.random_anisotropic_scale: dir_name = dir_name + '_aniso-scale' experiment_dir = experiment_dir.joinpath(dir_name) experiment_dir.mkdir(exist_ok=True) checkpoints_dir = experiment_dir.joinpath('checkpoints/') checkpoints_dir.mkdir(exist_ok=True) log_dir = experiment_dir.joinpath('logs/') log_dir.mkdir(exist_ok=True) # -------------------------------------------------------------------------- '''LOG''' # -------------------------------------------------------------------------- args = parse_args() logger = logging.getLogger("Model") logger.setLevel(logging.INFO) formatter = logging.Formatter( '%(asctime)s - %(name)s - %(levelname)s - %(message)s') file_handler = logging.FileHandler('%s/%s.txt' % (log_dir, args.model)) file_handler.setLevel(logging.INFO) file_handler.setFormatter(formatter) logger.addHandler(file_handler) log_string('PARAMETER ...') log_string(args) configure(log_dir) # tensorboard logdir log_string('OUTPUT DIR: %s' % experiment_dir) # -------------------------------------------------------------------------- '''DATA LOADERS''' # -------------------------------------------------------------------------- root = 'data/shapenetcore_partanno_segmentation_benchmark_v0_normal/' TRAIN_DATASET = PartNormalDataset(root=root, npoints=args.npoint, split='trainval', normal_channel=args.normal, k_shot=args.k_shot) trainDataLoader = torch.utils.data.DataLoader(TRAIN_DATASET, batch_size=args.batch_size, shuffle=True, num_workers=4) trainDataIterator = iter(trainDataLoader) TEST_DATASET = PartNormalDataset(root=root, npoints=args.npoint, split='test', normal_channel=args.normal) testDataLoader = torch.utils.data.DataLoader(TEST_DATASET, batch_size=args.batch_size, shuffle=False, num_workers=4) log_string("The number of training data is: %d" % len(TRAIN_DATASET)) log_string("The number of test data is: %d" % len(TEST_DATASET)) num_classes = 16 num_part = 50 if args.selfsup: log_string('Use self-supervision - alternate batches') if not args.retain_overlaps: log_string( '\tRemove overlaps between labeled and self-sup datasets') labeled_fns = list(itertools.chain(*TEST_DATASET.meta.values())) \ + list(itertools.chain(*TRAIN_DATASET.meta.values())) else: log_string('\tUse all files in self-sup dataset') labeled_fns = [] if args.ss_dataset == 'dummy': log_string( 'Using "dummy" self-supervision dataset (rest of labeled ShapeNetSeg)' ) SELFSUP_DATASET = SelfSupPartNormalDataset( root=root, npoints=args.npoint, split='trainval', normal_channel=args.normal, k_shot=args.n_cls_selfsup, labeled_fns=labeled_fns) elif args.ss_dataset == 'acd': log_string('Using "ACD" self-supervision dataset (ShapeNet Seg)') ACD_ROOT = args.ss_path SELFSUP_DATASET = ACDSelfSupDataset(root=ACD_ROOT, npoints=args.npoint, normal_channel=args.normal, k_shot=args.n_cls_selfsup, exclude_fns=labeled_fns, use_val=True) log_string('\t %d samples' % len(SELFSUP_DATASET)) selfsup_train_fns = list( itertools.chain(*SELFSUP_DATASET.meta.values())) log_string('Val dataset for self-sup') SELFSUP_VAL = ACDSelfSupDataset(root=ACD_ROOT, npoints=args.npoint, normal_channel=args.normal, class_choice='Airplane', k_shot=args.n_cls_selfsup, use_val=False, exclude_fns=selfsup_train_fns + labeled_fns) log_string('\t %d samples' % len(SELFSUP_VAL)) selfsupDataLoader = torch.utils.data.DataLoader( SELFSUP_DATASET, batch_size=args.batch_size, shuffle=True, num_workers=4) selfsupIterator = iter(selfsupDataLoader) selfsupValLoader = torch.utils.data.DataLoader( SELFSUP_VAL, batch_size=args.batch_size, shuffle=False, num_workers=4) log_string('Load ModelNet dataset for validation') DATA_PATH = 'data/modelnet40_normal_resampled/' MN_DATASET = ModelNetDataLoader(root=DATA_PATH, npoint=args.npoint, split='train', normal_channel=args.normal) modelnetLoader = torch.utils.data.DataLoader(MN_DATASET, batch_size=args.batch_size, shuffle=True, num_workers=4) # -------------------------------------------------------------------------- '''MODEL LOADING''' # -------------------------------------------------------------------------- MODEL = importlib.import_module(args.model) shutil.copy('models/%s.py' % args.model, str(experiment_dir)) shutil.copy('models/pointnet_util.py', str(experiment_dir)) if args.model == 'dgcnn': classifier = MODEL.get_model(num_part, normal_channel=args.normal, k=args.dgcnn_k).cuda() else: classifier = MODEL.get_model(num_part, normal_channel=args.normal).cuda() criterion = MODEL.get_loss().cuda() if args.selfsup: selfsupCriterion = MODEL.get_selfsup_loss(margin=args.margin).cuda() log_string("The number of self-sup data is: %d" % len(SELFSUP_DATASET)) def weights_init(m): classname = m.__class__.__name__ if classname.find('Conv2d') != -1: torch.nn.init.xavier_normal_(m.weight.data) torch.nn.init.constant_(m.bias.data, 0.0) elif classname.find('Linear') != -1: torch.nn.init.xavier_normal_(m.weight.data) torch.nn.init.constant_(m.bias.data, 0.0) try: checkpoint = torch.load( str(experiment_dir) + '/checkpoints/best_model.pth') start_epoch = checkpoint['epoch'] classifier.load_state_dict(checkpoint['model_state_dict']) log_string('Use pretrain model') except: log_string('No existing model, starting training from scratch...') start_epoch = 0 classifier = classifier.apply(weights_init) # -------------------------------------------------------------------------- '''OPTIMIZER SETTINGS''' # -------------------------------------------------------------------------- if args.optimizer == 'Adam': optimizer = torch.optim.Adam(classifier.parameters(), lr=args.learning_rate, betas=(0.9, 0.999), eps=1e-08, weight_decay=args.decay_rate) else: optimizer = torch.optim.SGD(classifier.parameters(), lr=args.learning_rate, momentum=0.9) def bn_momentum_adjust(m, momentum): if isinstance(m, torch.nn.BatchNorm2d) or isinstance( m, torch.nn.BatchNorm1d): m.momentum = momentum # LEARNING_RATE_CLIP = 1e-5 LEARNING_RATE_CLIP = args.lr_clip MOMENTUM_ORIGINAL = 0.1 MOMENTUM_DECAY = 0.5 MOMENTUM_DECAY_STEP = args.step_size if torch.cuda.device_count() > 1: print("Let's use", torch.cuda.device_count(), "GPUs!") classifier = nn.DataParallel(classifier) # -------------------------------------------------------------------------- '''TRAINING LOOP''' # -------------------------------------------------------------------------- best_val_loss = 99999 global_epoch = 0 for epoch in range(start_epoch, args.epoch): log_string('Epoch %d (%d/%s):' % (global_epoch + 1, epoch + 1, args.epoch)) '''Adjust learning rate and BN momentum''' lr = max( args.learning_rate * (args.lr_decay**(epoch // args.step_size)), LEARNING_RATE_CLIP) log_string('Learning rate:%f' % lr) for param_group in optimizer.param_groups: param_group['lr'] = lr mean_loss = [] momentum = MOMENTUM_ORIGINAL * (MOMENTUM_DECAY **(epoch // MOMENTUM_DECAY_STEP)) if momentum < 0.01: momentum = 0.01 print('BN momentum updated to: %f' % momentum) classifier = classifier.apply( lambda x: bn_momentum_adjust(x, momentum)) '''learning one epoch''' num_iters = len( selfsupDataLoader) # calc an epoch based on self-sup dataset for i in tqdm(list(range(num_iters)), total=num_iters, smoothing=0.9): '''applying self-supervised constrastive (pairwise) loss''' try: data_ss = next(selfsupIterator) except StopIteration: # reached end of this dataloader selfsupIterator = iter(selfsupDataLoader) data_ss = next(selfsupIterator) # DEBUG if DEBUG and i > 10: break points, label, target = data_ss # (points: bs x 3 x n_pts, label: bs x 1, target: bs x n_pts) points = points.data.numpy() points[:, :, 0:3] = provider.random_scale_point_cloud(points[:, :, 0:3]) points[:, :, 0:3] = provider.shift_point_cloud(points[:, :, 0:3]) if args.random_anisotropic_scale: points[:, :, 0:3] = provider.random_anisotropic_scale_point_cloud( points[:, :, 0:3], scale_low=0.8, scale_high=1.25) # pts = torch.Tensor(points) # pts = pts.transpose(2,1) # np.save(osp.join(experiment_dir, 'pts.npy'), pts.cpu().numpy()) if args.rotation_z: points[:, :, 0:3] = provider.rotate_point_cloud_y(points[:, :, 0:3]) if args.rotation_z_45: points[:, :, 0:3] = provider.rotate_point_cloud_y_pi4(points[:, :, 0:3]) points = torch.Tensor(points) points, label, target = points.float().cuda(), label.long().cuda( ), target.long().cuda() points = points.transpose(2, 1) # np.save(osp.join(experiment_dir, 'pts_z-rot.npy'), points.cpu().numpy()) # np.save(osp.join(experiment_dir, 'target.npy'), target.cpu().numpy()) # for self-sup category label is always unknown, so always zeros: category_label = torch.zeros([label.shape[0], 1, num_classes]).cuda() optimizer.zero_grad() classifier = classifier.train() _, _, feat = classifier(points, category_label) # feat: [bs x ndim x npts] ss_loss = selfsupCriterion(feat, target) * args.lmbda ss_loss.backward() optimizer.step() mean_loss.append(ss_loss.item()) log_value('selfsup_loss_iter', ss_loss.data, epoch * num_iters + i + 1) train_loss_epoch = np.mean(mean_loss) log_string('Self-sup loss is: %.5f' % train_loss_epoch) log_value('selfsup_loss_epoch', train_loss_epoch, epoch) # # # DEBUG: # with torch.no_grad(): # sa3_wt = classifier.sa3.mlp_convs[2].weight.mean() # log_string('SA3 avg wt is: %.5f' % sa3_wt.item()) # log_value('sa3_conv2_wt', sa3_wt.item(), epoch) '''validation after one epoch''' log_string('Validation: ACD on ShapeNet') with torch.no_grad(): total_val_loss = 0 for batch_id, (points, label, target) in tqdm(enumerate(selfsupValLoader), total=len(selfsupValLoader), smoothing=0.9): if DEBUG and i > 10: break cur_batch_size, NUM_POINT, _ = points.size() points, label, target = points.float().cuda(), label.long( ).cuda(), target.long().cuda() points = points.transpose(2, 1) category_label = torch.zeros([label.shape[0], 1, num_classes]).cuda() classifier = classifier.eval() _, _, feat = classifier(points, category_label) val_loss = selfsupCriterion(feat, target) total_val_loss += val_loss.data.cpu().item() avg_val_loss = total_val_loss / len(selfsupValLoader) log_value('selfsup_loss_val', avg_val_loss, epoch) '''(optional) validation on ModelNet40''' if args.modelnet_val: log_string('Validation: SVM on ModelNet40') with torch.no_grad(): log_string('Extract features on ModelNet40') if args.model == 'pointnet_part_seg': feat_train, label_train = extract_feats_pointnet( classifier, modelnetLoader, subset=0.5) elif args.model == 'pointnet2_part_seg_msg': feat_train, label_train = extract_feats(classifier, modelnetLoader, subset=0.5) else: raise ValueError log_string('Training data: %d samples, %d features' % feat_train.shape) start_time = time.time() log_string('Training SVM on ModelNet40') svm, best_C, best_score = cross_val_svm(feat_train, label_train, c_min=100, c_max=501, c_step=20, verbose=False) elapsed_time = time.time() - start_time log_string('ModelNet val Accuracy: %f (elapsed: %f seconds)' % (best_score, elapsed_time)) log_value('modelnet_val', best_score, epoch) # save every epoch if epoch % 5 == 0: savepath = str(checkpoints_dir) + ('/model_%03d.pth' % epoch) log_string('Saving model at %s' % savepath) state = { 'epoch': epoch, 'selfsup_loss': ss_loss.data, 'val_loss': avg_val_loss, 'model_state_dict': classifier.state_dict(), 'optimizer_state_dict': optimizer.state_dict(), } torch.save(state, savepath) log_string('Saved model.') # save best model if avg_val_loss < best_val_loss: best_val_loss = avg_val_loss savepath = str(checkpoints_dir) + '/best_model.pth' log_string('Saving best model at %s' % savepath) state = { 'epoch': epoch, 'selfsup_loss': ss_loss.data, 'val_loss': avg_val_loss, 'model_state_dict': classifier.state_dict(), 'optimizer_state_dict': optimizer.state_dict(), } torch.save(state, savepath) log_string('Saved model.') log_value('train_lr', lr, epoch) log_value('train_bn_momentum', momentum, epoch) log_string('Epoch %d Self-sup train loss: %f Val loss: %f ' % (epoch + 1, train_loss_epoch, avg_val_loss)) global_epoch += 1
def train(args): experiment_dir = mkdir('./experiment/') checkpoints_dir = mkdir('./experiment/clf/%s/' % (args.model_name)) train_data, train_label, test_data, test_label = load_data( 'experiment/data/modelnet40_ply_hdf5_2048/') trainDataset = ModelNetDataLoader(train_data, train_label, data_augmentation=args.augment) trainDataLoader = DataLoader(trainDataset, batch_size=args.batch_size, shuffle=True) testDataset = ModelNetDataLoader(test_data, test_label) testDataLoader = torch.utils.data.DataLoader(testDataset, batch_size=args.batch_size, shuffle=False) log.info('Building Model', args.model_name) if args.model_name == 'pointnet': num_class = 40 model = PointNetCls(num_class, args.feature_transform).cuda() else: model = PointNet2ClsMsg().cuda() torch.backends.cudnn.benchmark = True model = torch.nn.DataParallel(model).cuda() log.debug('Using gpu:', args.gpu) if args.pretrain is not None: log.info('Use pretrain model...') state_dict = torch.load(args.pretrain) model.load_state_dict(state_dict) init_epoch = int(args.pretrain[:-4].split('-')[-1]) log.info('start epoch from', init_epoch) else: log.info('Training from scratch') init_epoch = 0 if args.optimizer == 'SGD': optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9) elif args.optimizer == 'Adam': optimizer = torch.optim.Adam(model.parameters(), lr=args.learning_rate, betas=(0.9, 0.999), eps=1e-08, weight_decay=args.decay_rate) scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=20, gamma=0.5) LEARNING_RATE_CLIP = 1e-5 global_epoch = 0 global_step = 0 best_tst_accuracy = 0.0 log.info('Start training...') for epoch in range(init_epoch, args.epoch): scheduler.step() lr = max(optimizer.param_groups[0]['lr'], LEARNING_RATE_CLIP) log.debug(job='clf', model=args.model_name, gpu=args.gpu, epoch='%d/%s' % (epoch, args.epoch), lr=lr) for param_group in optimizer.param_groups: param_group['lr'] = lr for batch_id, data in tqdm(enumerate(trainDataLoader, 0), total=len(trainDataLoader), smoothing=0.9): points, target = data target = target[:, 0] points = points.transpose(2, 1) points, target = points.cuda(), target.cuda() optimizer.zero_grad() model = model.train() pred, trans_feat = model(points) loss = F.nll_loss(pred, target.long()) if args.feature_transform and args.model_name == 'pointnet': loss += feature_transform_reguliarzer(trans_feat) * 0.001 loss.backward() optimizer.step() global_step += 1 log.debug('clear cuda cache') torch.cuda.empty_cache() acc = test_clf(model, testDataLoader) log.info(loss='%.5f' % (loss.data)) log.info(Test_Accuracy='%.5f' % acc) if acc >= best_tst_accuracy: best_tst_accuracy = acc fn_pth = 'clf-%s-%.5f-%04d.pth' % (args.model_name, acc, epoch) log.debug('Saving model....', fn_pth) torch.save(model.state_dict(), os.path.join(checkpoints_dir, fn_pth)) global_epoch += 1 log.info(Best_Accuracy=best_tst_accuracy) log.info('End of training...')
def main(args): '''HYPER PARAMETER''' os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu # datapath = './data/ModelNet/' datapath = './data/modelnet40_ply_hdf5_2048/' if args.rotation is not None: ROTATION = (int(args.rotation[0:2]), int(args.rotation[3:5])) else: ROTATION = None '''CREATE DIR''' experiment_dir = Path('./experiment/') experiment_dir.mkdir(exist_ok=True) checkpoints_dir = Path('./experiment/checkpoints/') checkpoints_dir.mkdir(exist_ok=True) log_dir = Path('./experiment/logs/') log_dir.mkdir(exist_ok=True) '''LOG''' args = parse_args() logger = logging.getLogger("PointNet2") logger.setLevel(logging.INFO) formatter = logging.Formatter( '%(asctime)s - %(name)s - %(levelname)s - %(message)s') file_handler = logging.FileHandler( './experiment/logs/test_%s_' % args.model_name + str(datetime.datetime.now().strftime('%Y-%m-%d %H-%M')) + '.txt') file_handler.setLevel(logging.INFO) file_handler.setFormatter(formatter) logger.addHandler(file_handler) logger.info( '---------------------------------------------------Test---------------------------------------------------' ) logger.info('PARAMETER ...') logger.info(args) '''DATA LOADING''' logger.info('Load dataset ...') train_data, train_label, test_data, test_label = load_data( datapath, classification=True) logger.info("The number of training data is: %d", train_data.shape[0]) logger.info("The number of test data is: %d", test_data.shape[0]) trainDataset = ModelNetDataLoader(train_data, train_label, rotation=ROTATION) if ROTATION is not None: print('The range of training rotation is', ROTATION) testDataset = ModelNetDataLoader(test_data, test_label, rotation=ROTATION) trainDataLoader = torch.utils.data.DataLoader(trainDataset, batch_size=args.batchsize, shuffle=True) testDataLoader = torch.utils.data.DataLoader(testDataset, batch_size=args.batchsize, shuffle=False) # 不打乱 '''MODEL LOADING''' num_class = 40 ###################### PointNetCls ###################### classifier = PointNetCls(num_class, args.feature_transform).cuda( ) if args.model_name == 'pointnet' else PointNet2ClsMsg().cuda() if args.pretrain is not None: print('Use pretrain model...') logger.info('Use pretrain model') checkpoint = torch.load(args.pretrain) start_epoch = checkpoint['epoch'] classifier.load_state_dict(checkpoint['model_state_dict']) else: print('Please Input the pretrained model ***.pth') return if args.optimizer == 'SGD': optimizer = torch.optim.SGD(classifier.parameters(), lr=0.01, momentum=0.9) elif args.optimizer == 'Adam': optimizer = torch.optim.Adam(classifier.parameters(), lr=args.learning_rate, betas=(0.9, 0.999), eps=1e-08, weight_decay=args.decay_rate) scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=20, gamma=0.5) global_epoch = 0 global_step = 0 best_tst_accuracy = 0.0 blue = lambda x: '\033[94m' + x + '\033[0m' '''TestING''' logger.info('Start testing...') scheduler.step() acc = test(classifier.eval(), testDataLoader) print('\r Test %s: %f' % (blue('Accuracy'), acc)) logger.info('Test Accuracy: %f', acc) logger.info('End of testing...')
def main(args): def log_string(str): logger.info(str) print(str) '''HYPER PARAMETER''' if not args.reduced_computation: os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu num_workers = 8 else: num_workers = 0 '''CREATE DIR''' timestr = str(datetime.datetime.now().strftime('%Y-%m-%d_%H-%M')) experiment_dir = Path(args.experiment_dir) experiment_dir.mkdir(exist_ok=True) experiment_dir = experiment_dir.joinpath('classification') experiment_dir.mkdir(exist_ok=True) if args.log_dir is None: experiment_dir = experiment_dir.joinpath(timestr) else: experiment_dir = experiment_dir.joinpath(args.log_dir) experiment_dir.mkdir(exist_ok=True) checkpoints_dir = experiment_dir.joinpath('checkpoints/') checkpoints_dir.mkdir(exist_ok=True) log_dir = experiment_dir.joinpath('logs/') log_dir.mkdir(exist_ok=True) tensorboard_dir = experiment_dir.joinpath('tensorboard/') tensorboard_dir.mkdir(exist_ok=True) '''LOG''' args = parse_args() logger = logging.getLogger("Model") logger.setLevel(logging.INFO) formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s') file_handler = logging.FileHandler('%s/%s.txt' % (log_dir, args.model)) file_handler.setLevel(logging.INFO) file_handler.setFormatter(formatter) logger.addHandler(file_handler) log_string('PARAMETER ...') log_string(args) '''TENSORBOARD''' train_writer = SummaryWriter(tensorboard_dir.joinpath("train")) val_writer = SummaryWriter(tensorboard_dir.joinpath("validation")) '''DATA LOADING''' log_string('Load dataset ...') DATA_PATH = args.data_path class_in_filename = False if args.data_extension == ".npy" else True TRAIN_DATASET = ModelNetDataLoader(root=DATA_PATH, split_name=args.split_name, extension=args.data_extension, npoint=args.num_point, split='train', normal_channel=args.normal, class_in_filename=class_in_filename, uniform=args.uniform, voxel_size=args.reduced_resolution_voxel_size) TEST_DATASET = ModelNetDataLoader(root=DATA_PATH, split_name=args.split_name, extension=args.data_extension, npoint=args.num_point, split='validation', normal_channel=args.normal, class_in_filename=class_in_filename, uniform=args.uniform, voxel_size=args.reduced_resolution_voxel_size) trainDataLoader = torch.utils.data.DataLoader(TRAIN_DATASET, batch_size=args.batch_size, shuffle=True, num_workers=num_workers) testDataLoader = torch.utils.data.DataLoader(TEST_DATASET, batch_size=args.batch_size, shuffle=False, num_workers=num_workers) '''MODEL LOADING''' num_class = args.num_classes MODEL = importlib.import_module(args.model) shutil.copy('./models/%s.py' % args.model, str(experiment_dir)) shutil.copy('./models/pointnet_util.py', str(experiment_dir)) classifier = MODEL.get_model(num_class, normal_channel=args.normal) if not args.reduced_computation: classifier = classifier.cuda() criterion = MODEL.get_loss() if not args.reduced_computation: criterion = criterion.cuda() try: checkpoint = torch.load(str(experiment_dir) + '/checkpoints/best_model.pth') start_epoch = checkpoint['epoch'] classifier.load_state_dict(checkpoint['model_state_dict']) log_string('Use pretrain model') except: log_string('No existing model, starting training from scratch...') start_epoch = 0 if args.optimizer == 'Adam': optimizer = torch.optim.Adam( classifier.parameters(), lr=args.learning_rate, betas=(0.9, 0.999), eps=1e-08, weight_decay=args.decay_rate ) else: optimizer = torch.optim.SGD(classifier.parameters(), lr=0.01, momentum=0.9) scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.9, patience=10, min_lr=0.000001) global_epoch = 0 global_step = 0 best_instance_acc = 0.0 best_class_acc = 0.0 mean_correct = [] '''TRANING''' logger.info('Start training...') for epoch in range(start_epoch,args.epoch): log_string('Epoch %d (%d/%s):' % (global_epoch + 1, epoch + 1, args.epoch)) mean_correct = [] batch_tqdm = tqdm(enumerate(trainDataLoader, 0), total=len(trainDataLoader), smoothing=0.9) total_loss = 0 predictions_likelihood_tot = torch.zeros([len(trainDataLoader.dataset), num_class]) for batch_id, data in batch_tqdm: points, target = data points = points.data.numpy() if args.augment: points = provider.random_point_dropout(points) points[:,:, 0:3] = provider.random_scale_point_cloud(points[:,:, 0:3]) points[:,:, 0:3] = provider.shift_point_cloud(points[:,:, 0:3]) points = torch.Tensor(points) target = target[:, 0] points = points.transpose(2, 1) if not args.reduced_computation: points, target = points.cuda(), target.cuda() optimizer.zero_grad() classifier = classifier.train() pred, trans_feat = classifier(points) loss = criterion(pred, target.long(), trans_feat) pred_choice = pred.data.max(1)[1] correct = pred_choice.eq(target.long().data).cpu().sum() mean_correct.append(correct.item() / float(points.size()[0])) loss.backward() optimizer.step() global_step += 1 total_loss += loss mean_loss = total_loss / (batch_id + 1) batch_tqdm.set_description(f"loss {mean_loss}, batch ({batch_id}/{len(trainDataLoader)})") preds_likelihood = torch.exp(pred) predictions_likelihood_tot[batch_id*trainDataLoader.batch_size:(batch_id+1)*trainDataLoader.batch_size] = preds_likelihood train_instance_acc = np.mean(mean_correct) log_string('Train Instance Accuracy: %f' % train_instance_acc) train_writer.add_scalar('Loss', mean_loss, epoch) train_writer.add_scalar('Accuracy', train_instance_acc, epoch) for cls in range(num_class): train_writer.add_histogram(f"class_{cls}", predictions_likelihood_tot[:, cls], epoch) with torch.no_grad(): instance_acc, class_acc, val_loss = test(classifier.eval(), testDataLoader, criterion, num_class=num_class) scheduler.step(val_loss) val_writer.add_scalar('Loss', val_loss, epoch) val_writer.add_scalar('Accuracy', instance_acc, epoch) val_writer.add_scalar('Class_Accuracy', class_acc, epoch) if (instance_acc >= best_instance_acc): best_instance_acc = instance_acc best_epoch = epoch + 1 if (class_acc >= best_class_acc): best_class_acc = class_acc log_string('Test Instance Accuracy: %f, Class Accuracy: %f'% (instance_acc, class_acc)) log_string('Best Instance Accuracy: %f, Class Accuracy: %f'% (best_instance_acc, best_class_acc)) if (instance_acc >= best_instance_acc): logger.info('Save model...') savepath = str(checkpoints_dir) + '/best_model.pth' log_string('Saving at %s'% savepath) state = { 'epoch': best_epoch, 'instance_acc': instance_acc, 'class_acc': class_acc, 'model_state_dict': classifier.state_dict(), 'optimizer_state_dict': optimizer.state_dict(), } torch.save(state, savepath) global_epoch += 1 logger.info('End of training...')
def main(args): """HYPER PARAMETER""" os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu datapath = './data/ModelNet/' '''CREATE DIR''' experiment_dir = Path('./eval_experiment/') experiment_dir.mkdir(exist_ok=True) file_dir = Path(str(experiment_dir) + '/%sModelNet40-' % args.model_name + str( datetime.datetime.now().strftime('%Y-%m-%d_%H-%M'))) file_dir.mkdir(exist_ok=True) checkpoints_dir = file_dir.joinpath('checkpoints/') checkpoints_dir.mkdir(exist_ok=True) os.system('cp %s %s' % (args.checkpoint, checkpoints_dir)) log_dir = file_dir.joinpath('logs/') log_dir.mkdir(exist_ok=True) '''LOG''' args = parse_args() logger = logging.getLogger(args.model_name) logger.setLevel(logging.INFO) formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s') file_handler = logging.FileHandler(str(log_dir) + 'eval_%s_cls.txt' % args.model_name) file_handler.setLevel(logging.INFO) file_handler.setFormatter(formatter) logger.addHandler(file_handler) logger.info( '---------------------------------------------------EVAL---------------------------------------------------') logger.info('PARAMETER ...') logger.info(args) '''DATA LOADING''' logger.info('Load dataset ...') train_data, train_label, test_data, test_label = load_data(datapath, classification=True) logger.info("The number of training data is: %d", train_data.shape[0]) logger.info("The number of test data is: %d", test_data.shape[0]) test_dataset = ModelNetDataLoader(test_data, test_label) test_data_loader = torch.utils.data.DataLoader(test_dataset, batch_size=args.batchsize, shuffle=False) '''MODEL LOADING''' num_class = 40 classifier = PointConvClsSsg(num_class).cuda() if args.checkpoint is not None: print('Load CheckPoint...') logger.info('Load CheckPoint') checkpoint = torch.load(args.checkpoint) start_epoch = checkpoint['epoch'] classifier.load_state_dict(checkpoint['model_state_dict']) else: print('Please load Checkpoint to eval...') start_epoch = 0 sys.exit(0) blue = lambda x: '\033[94m' + x + '\033[0m' '''EVAL''' logger.info('Start evaluating...') print('Start evaluating...') total_correct = 0 total_seen = 0 for batch_id, data in tqdm(enumerate(test_data_loader, 0), total=len(test_data_loader), smoothing=0.9): pointcloud, target = data target = target[:, 0] # import ipdb; ipdb.set_trace() pred_view = torch.zeros(pointcloud.shape[0], num_class).cuda() for _ in range(args.num_view): pointcloud = generate_new_view(pointcloud) # import ipdb; ipdb.set_trace() # points = torch.from_numpy(pointcloud).permute(0, 2, 1) points = pointcloud.permute(0, 2, 1) points, target = points.cuda(), target.cuda() classifier = classifier.eval() with torch.no_grad(): pred = classifier(points) pred_view += pred pred_choice = pred_view.data.max(1)[1] correct = pred_choice.eq(target.long().data).cpu().sum() total_correct += correct.item() total_seen += float(points.size()[0]) accuracy = total_correct / total_seen print('Total Accuracy: %f' % accuracy) logger.info('Total Accuracy: %f' % accuracy) logger.info('End of evaluation...')
def main(args): def log_string(str): logger.info(str) print(str) '''HYPER PARAMETER''' os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu '''CREATE DIR''' timestr = str(datetime.datetime.now().strftime('%Y-%m-%d_%H-%M')) experiment_dir = Path('./log/') experiment_dir.mkdir(exist_ok=True) dataset_name = args.dataset_name experiment_dir = experiment_dir.joinpath( 'classification_{}'.format(dataset_name)) experiment_dir.mkdir(exist_ok=True) if args.log_dir is None: experiment_dir = experiment_dir.joinpath(timestr) else: experiment_dir = experiment_dir.joinpath(args.log_dir) experiment_dir.mkdir(exist_ok=True) checkpoints_dir = experiment_dir.joinpath('checkpoints/') checkpoints_dir.mkdir(exist_ok=True) log_dir = experiment_dir.joinpath('logs/') log_dir.mkdir(exist_ok=True) '''LOG''' args = parse_args() logger = logging.getLogger("Model") logger.setLevel(logging.INFO) formatter = logging.Formatter( '%(asctime)s - %(name)s - %(levelname)s - %(message)s') file_handler = logging.FileHandler('%s/%s.txt' % (log_dir, args.model)) file_handler.setLevel(logging.INFO) file_handler.setFormatter(formatter) logger.addHandler(file_handler) log_string('PARAMETER ...') log_string(args) '''TENSORBOARD LOG''' writer = SummaryWriter() '''DATA LOADING''' log_string('Load dataset ...') DATA_PATH = os.path.join(ROOT_DIR, 'data', dataset_name) print("loading dataset from {}".format(dataset_name)) if 'modelnet' in dataset_name: TRAIN_DATASET = ModelNetDataLoader(DATA_PATH, split='train', normal_channel=args.normal) TEST_DATASET = ModelNetDataLoader(DATA_PATH, split='test', normal_channel=args.normal) num_class = 40 else: print(DATA_PATH) TRAIN_DATASET = ReplicaDataLoader(DATA_PATH, split='train', uniform=True, normal_channel=False, rot_transform=True) TEST_DATASET = ReplicaDataLoader(DATA_PATH, split='test', uniform=True, normal_channel=False, rot_transform=False) num_class = 31 trainDataLoader = torch.utils.data.DataLoader(TRAIN_DATASET, batch_size=args.batch_size, shuffle=True, num_workers=6) testDataLoader = torch.utils.data.DataLoader(TEST_DATASET, batch_size=args.batch_size, shuffle=False, num_workers=6) '''MODEL LOADING''' print("Number of classes are {:d}".format(num_class)) MODEL = importlib.import_module(args.model) shutil.copy('./models/%s.py' % args.model, str(experiment_dir)) shutil.copy('./models/pointnet_util.py', str(experiment_dir)) print("Obtain GPU device ") train_GPU = True device = torch.device("cuda" if ( torch.cuda.is_available() and train_GPU) else "cpu") print(device) print("Load the network to the device ") classifier = MODEL.get_model(num_class, normal_channel=args.normal).to(device) print("Load the loss to the device ") criterion = MODEL.get_loss().to(device) if os.path.exists((str(experiment_dir) + '/checkpoints/best_model.pth')): checkpoint = torch.load( str(experiment_dir) + '/checkpoints/best_model.pth') start_epoch = checkpoint['epoch'] classifier.load_state_dict(checkpoint['model_state_dict']) # strict set to false to allow using the model trained with modelnet else: start_epoch = 0 if dataset_name == 'replica': log_string('Use pretrain model of Model net') # double check again if there is pretrained modelnet model checkpoint = torch.load( str(experiment_dir).replace("replica", 'modelnet40_normal_resampled') + '/checkpoints/best_model.pth') classifier = MODEL.get_model(40, normal_channel=args.normal).to(device) classifier.load_state_dict(checkpoint['model_state_dict']) classifier.fc3 = nn.Linear(256, num_class).to(device) print(classifier) else: log_string('No existing model, starting training from scratch...') if args.optimizer == 'Adam': print("Using Adam opimizer ") optimizer = torch.optim.Adam(classifier.parameters(), lr=args.learning_rate, betas=(0.9, 0.999), eps=1e-08, weight_decay=args.decay_rate) else: optimizer = torch.optim.SGD(classifier.parameters(), lr=0.01, momentum=0.9) scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=20, gamma=0.7) global_epoch = 0 global_step = 0 best_instance_acc = 0.0 best_class_acc = 0.0 mean_correct = [] '''TRANING''' logger.info('Start training...') for epoch in range(start_epoch, args.epoch): loss_array = np.zeros((len(trainDataLoader), 1)) log_string('Epoch %d (%d/%s):' % (global_epoch + 1, epoch + 1, args.epoch)) classifier.train() # setting the model to train mode print("Clear GPU cache ...") torch.cuda.empty_cache() scheduler.step() for batch_id, data in tqdm(enumerate(trainDataLoader, 0), total=len(trainDataLoader), smoothing=0.9): points, target = data points = points.data.numpy() points = provider.random_point_dropout(points) points[:, :, 0:3] = provider.random_scale_point_cloud(points[:, :, 0:3]) points[:, :, 0:3] = provider.shift_point_cloud(points[:, :, 0:3]) points = torch.Tensor(points) target = target[:, 0] points = points.transpose(2, 1) points, target = points.to(device), target.to(device) optimizer.zero_grad() pred, trans_feat = classifier( points) ### This is the part of the runtime error: loss = criterion(pred, target.long(), trans_feat) loss_array[batch_id] = loss.cpu().detach().numpy() pred_choice = pred.data.max(1)[1] correct = pred_choice.eq(target.long().data).cpu().sum() mean_correct.append(correct.item() / float(points.size()[0])) loss.backward() optimizer.step() global_step += 1 train_instance_acc = np.mean(mean_correct) log_string('Train Instance Accuracy: %f' % train_instance_acc) avg_loss = np.mean(loss_array[:]) writer.add_scalar("Loss/train", avg_loss, epoch) ## This is for validation with torch.no_grad(): instance_acc, class_acc = test(classifier.eval(), testDataLoader, device, num_class) writer.add_scalar("ClassAccuracy/test", class_acc, epoch) writer.add_scalar("InstanceAccuracy/test", instance_acc, epoch) if (instance_acc >= best_instance_acc): best_instance_acc = instance_acc best_epoch = epoch + 1 if (class_acc >= best_class_acc): best_class_acc = class_acc log_string('Test Instance Accuracy: %f, Class Accuracy: %f' % (instance_acc, class_acc)) log_string('Best Instance Accuracy: %f, Class Accuracy: %f' % (best_instance_acc, best_class_acc)) if (instance_acc >= best_instance_acc): logger.info('Save model...') savepath = str(checkpoints_dir) + '/best_model.pth'.format( epoch) log_string('Saving at %s' % savepath) state = { 'epoch': best_epoch, 'instance_acc': instance_acc, 'class_acc': class_acc, 'model_state_dict': classifier.state_dict(), 'optimizer_state_dict': optimizer.state_dict(), } torch.save(state, savepath) global_epoch += 1 logger.info('End of training...') writer.flush() writer.close()
def main(args): def log_string(str): logger.info(str) print(str) '''HYPER PARAMETER''' os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu '''CREATE DIR''' experiment_dir = osp.join(args.log_dir, 'ModelNet40-eval') experiment_dir = experiment_dir + '_' + str(args.num_point) if args.sqrt: experiment_dir = experiment_dir + '_do-sqrt' if args.do_sa3: experiment_dir = experiment_dir + '_sa3-feats' if args.svm_jitter: experiment_dir = experiment_dir + '_svm-jitter' args.batch_size = (args.batch_size // 8) # 8x augmentation if args.random_feats: experiment_dir = experiment_dir + '_random-feats' if args.ckpt is not None: experiment_dir = experiment_dir + '_' + osp.splitext(args.ckpt)[0] os.makedirs(experiment_dir, exist_ok=True) '''LOG''' logger = logging.getLogger("Model") logger.setLevel(logging.INFO) formatter = logging.Formatter( '%(asctime)s - %(name)s - %(levelname)s - %(message)s') file_handler = logging.FileHandler('%s/eval.txt' % experiment_dir) file_handler.setLevel(logging.INFO) file_handler.setFormatter(formatter) logger.addHandler(file_handler) log_string('PARAMETER ...') log_string(args) log_string('Experiment dir: %s' % experiment_dir) '''DATA LOADING''' log_string('Load dataset ...') DATA_PATH = 'data/modelnet40_normal_resampled/' TRAIN_DATASET = ModelNetDataLoader(root=DATA_PATH, npoint=args.num_point, split='train', normal_channel=args.normal) TEST_DATASET = ModelNetDataLoader(root=DATA_PATH, npoint=args.num_point, split='test', normal_channel=args.normal) trainDataLoader = torch.utils.data.DataLoader(TRAIN_DATASET, batch_size=args.batch_size, shuffle=False, num_workers=4) testDataLoader = torch.utils.data.DataLoader(TEST_DATASET, batch_size=args.batch_size, shuffle=False, num_workers=4) if DEBUG: # ShapeNet training data shapenet_root = 'data/shapenetcore_partanno_segmentation_benchmark_v0_normal/' SHAPENET_DATASET = PartNormalDataset(root=shapenet_root, npoints=args.num_point, split='trainval', normal_channel=args.normal) shapenetDataLoader = torch.utils.data.DataLoader( SHAPENET_DATASET, batch_size=args.batch_size, shuffle=False, num_workers=4) ACD_ROOT = '/srv/data2/mgadelha/ShapeNetACD/' SELFSUP_DATASET = ACDSelfSupDataset(root=ACD_ROOT, npoints=args.num_point, normal_channel=args.normal) selfsupDataLoader = torch.utils.data.DataLoader( SELFSUP_DATASET, batch_size=args.batch_size, shuffle=False, num_workers=4) '''MODEL LOADING''' shapenet_num_class = 50 # model_name = args.model MODEL = importlib.import_module(model_name) model = MODEL.get_model(shapenet_num_class, normal_channel=False).cuda() if not args.random_feats: log_string('Load ACD pre-trained model: %s' % args.log_dir) if args.ckpt is None: checkpoint = torch.load( str(args.log_dir) + '/checkpoints/best_model.pth') else: checkpoint = torch.load( str(args.log_dir) + '/checkpoints/' + args.ckpt) try: DATA_PARALLEL = False model.load_state_dict(checkpoint['model_state_dict']) except: DATA_PARALLEL = True model = nn.DataParallel(model) model.load_state_dict(checkpoint['model_state_dict']) model = model.module else: log_string('Using randomly initialized %s as feature extractor' % model_name) # Extract features and save if not osp.exists(osp.join(experiment_dir, 'train-feats.npy')) or \ not osp.exists(osp.join(experiment_dir, 'train-labels.txt')): log_string('Extract features ...') if args.model == 'pointnet_part_seg': feat_train, label_train = extract_feats_pointnet( model, trainDataLoader, do_sqrt=args.sqrt, do_global=args.do_sa3) feat_test, label_test = extract_feats_pointnet( model, testDataLoader, do_sqrt=args.sqrt, do_global=args.do_sa3) elif args.model == 'pointnet2_part_seg_msg': feat_train, label_train = extract_feats( model, trainDataLoader, do_sqrt=args.sqrt, do_sa3=args.do_sa3, do_svm_jitter=args.svm_jitter) feat_test, label_test = extract_feats( model, testDataLoader, do_sqrt=args.sqrt, do_sa3=args.do_sa3, do_svm_jitter=args.svm_jitter) elif args.model == 'dgcnn': pass # feat_train, label_train = extract_feats_dgcnn(model, trainDataLoader, # do_sqrt=args.sqrt) # feat_test, label_test = extract_feats_dgcnn(model, testDataLoader, # do_sqrt=args.sqrt) elif args.model == 'dgcnn_seg': feat_train, label_train = extract_feats_dgcnn(model, trainDataLoader, do_sqrt=args.sqrt) feat_test, label_test = extract_feats_dgcnn(model, testDataLoader, do_sqrt=args.sqrt) else: raise ValueError np.save(osp.join(experiment_dir, 'train-feats.npy'), feat_train) np.savetxt(osp.join(experiment_dir, 'train-labels.txt'), label_train) np.save(osp.join(experiment_dir, 'test-feats.npy'), feat_test) np.savetxt(osp.join(experiment_dir, 'test-labels.txt'), label_test) else: log_string('Loading pre-trained features') feat_train = np.load(osp.join(experiment_dir, 'train-feats.npy')) label_train = np.loadtxt(osp.join(experiment_dir, 'train-labels.txt')) feat_test = np.load(osp.join(experiment_dir, 'test-feats.npy')) label_test = np.loadtxt(osp.join(experiment_dir, 'test-labels.txt')) # Train linear SVM (one-vs-rest) on features # Train+test SVM on validation *or* test set log_string('Training linear SVM ...') if args.val_svm: log_string('Total data: %d samples, %d features' % feat_train.shape) val_acc, _, _ = train_val_svm(feat_train, label_train, svm_c=args.svm_c) log_string('Validation Accuracy: %f' % val_acc) else: # SVM training on *all* training data log_string('Training data: %d samples, %d features' % feat_train.shape) t_0 = time.time() if args.cross_val_svm: classifier, best_C, best_score = cross_val_svm( feat_train, label_train) else: classifier = LinearSVC(random_state=123, multi_class='ovr', C=args.svm_c, dual=False) classifier.fit(feat_train, label_train) train_acc = classifier.score(feat_train, label_train) log_string('Train Accuracy: %f' % train_acc) t_1 = time.time() log_string('Time elapsed: %f' % (t_1 - t_0)) # test performance test_acc = classifier.score(feat_test, label_test) log_string('Test Accuracy: %f' % test_acc)
def get_acc(): batch_size = 8 config = yaml.load(open("/content/PointNet-BYOL/config/config.yaml", "r"), Loader=yaml.FullLoader) #这里normal_channel一定要改成True,不然channel会变成3,无法与6匹配 TRAIN_DATASET = ModelNetDataLoader( root='data/modelnet40_normal_resampled/', npoint=1024, split='train', normal_channel=True) TEST_DATASET = ModelNetDataLoader(root='data/modelnet40_normal_resampled/', npoint=1024, split='test', normal_channel=True) print("Input shape:", len(TRAIN_DATASET)) train_loader = torch.utils.data.DataLoader(TRAIN_DATASET, batch_size=16, shuffle=True, num_workers=12) test_loader = torch.utils.data.DataLoader(TEST_DATASET, batch_size=16, shuffle=False, num_workers=12) device = 'cuda' if torch.cuda.is_available( ) else 'cpu' #'cuda' if torch.cuda.is_available() else 'cpu' encoder = get_model(num_class=40, normal_channel=True) # output_feature_dim = encoder.projetion.net[0].in_features# # load pre-trained parameters load_params = torch.load( os.path.join('/content/PointNet-BYOL/checkpoints/model.pth'), map_location=torch.device(torch.device(device))) if 'online_network_state_dict' in load_params: encoder.load_state_dict(load_params['online_network_state_dict']) print("Parameters successfully loaded.") # remove the projection head encoder = encoder.to(device) encoder.eval() # x_train, y_train = get_features_from_encoder(encoder, train_loader) # x_test, y_test = get_features_from_encoder(encoder, test_loader) # x_train = torch.mean(x_train, dim=[2, 3]) # x_test = torch.mean(x_test, dim=[2, 3]) # print("Training data shape:", x_train.shape, y_train.shape) # print("Testing data shape:", x_test.shape, y_test.shape) # x_train=np.array(x_train) # scaler = preprocessing.StandardScaler() # scaler.fit(x_train) # x_train = scaler.transform(x_train).astype(np.float32) # x_test = scaler.transform(x_test).astype(np.float32) # train_loader, test_loader = create_data_loaders_from_arrays(torch.tensor([item.cpu().detach().numpy() for item in x_train]).cuda(), \ # y_train, torch.from_numpy(x_test), y_test) # train_loader, test_loader = create_data_loaders_from_arrays(torch.from_numpy(x_train), \ # y_train, torch.from_numpy(x_test), y_test) criterion = pointnet2_cls_msg_concat.get_loss().cuda() classifier = pointnet2_cls_msg_concat.get_model( num_class=40, normal_channel=True).cuda() # optimizer = torch.optim.SGD(classifier.parameters(), lr=0.01, momentum=0.9) optimizer = torch.optim.Adam(classifier.parameters(), lr=0.001, betas=(0.9, 0.999), eps=1e-08, weight_decay=1e-4) eval_every_n_epochs = 1 scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=20, gamma=0.7) best_instance_acc = 0.0 best_class_acc = 0.0 mean_correct = [] try: checkpoint = torch.load( '/content/PointNet-BYOL/checkpoints/best_model.pth') classifier.load_state_dict(checkpoint['model_state_dict']) optimizer.load_state_dict(checkpoint['optimizer_state_dict']) print('Use pretrain model') except: print('No existing model, starting training from scratch...') for epoch in range(20): print('Epoch %d ' % (epoch + 1)) scheduler.step() for batch_id, data in tqdm(enumerate(train_loader, 0), total=len(train_loader), smoothing=0.9): points, target = data points = points.data.numpy() points = provider.random_point_dropout(points) points[:, :, 0:3] = provider.random_scale_point_cloud(points[:, :, 0:3]) points[:, :, 0:3] = provider.shift_point_cloud(points[:, :, 0:3]) points = torch.Tensor(points) target = target[:, 0] points = points.transpose(2, 1) points, target = points.cuda(), target.cuda() optimizer.zero_grad() classifier = classifier.train() feature_vector, cls = encoder(points) pred, cls = classifier(points, feature_vector) loss = criterion(pred, target.long(), target) pred_choice = pred.data.max(1)[1] correct = pred_choice.eq(target.long().data).cpu().sum() mean_correct.append(correct.item() / float(points.size()[0])) loss.backward() optimizer.step() train_instance_acc = np.mean(mean_correct) print('Train Instance Accuracy: %f' % train_instance_acc) with torch.no_grad(): instance_acc, class_acc = test(classifier.eval(), encoder.eval(), test_loader) if (instance_acc >= best_instance_acc): best_instance_acc = instance_acc best_epoch = epoch + 1 if (class_acc >= best_class_acc): best_class_acc = class_acc print('Test Instance Accuracy: %f, Class Accuracy: %f' % (instance_acc, class_acc)) print('Best Instance Accuracy: %f, Class Accuracy: %f' % (best_instance_acc, best_class_acc)) if (instance_acc >= best_instance_acc): print('Save model...') savepath = 'checkpoints' + '/best_model.pth' print('Saving at %s' % savepath) state = { 'epoch': best_epoch, 'instance_acc': instance_acc, 'class_acc': class_acc, 'model_state_dict': classifier.state_dict(), 'optimizer_state_dict': optimizer.state_dict(), } torch.save(state, savepath) print('End of training...') # train = torch.utils.data.TensorDataset(pred1, target) # train_loader = torch.utils.data.DataLoader(train, batch_size=96, shuffle=True) # for epoch in range(10): # train_acc = [] # for x, y in train_loader: # x = x.to(device) # y = y.to(device) # # zero the parameter gradients # optimizer.zero_grad() # classifier = classifier.train() # pred = classifier(x) # predictions = torch.argmax(pred, dim=1) # loss = criterion(pred, y.long(),y) # loss.backward(retain_graph=True) # optimizer.step() # if epoch % eval_every_n_epochs == 0: # train_total,total = 0,0 # train_correct,correct = 0,0 # for x, y in train_loader: # x = x.to(device) # y = y.to(device) # classifier = classifier.train() # pred = classifier(x) # predictions = torch.argmax(pred, dim=1) # train_total += y.size(0) # train_correct += (predictions == y).sum().item() # for x, y in test_loader: # x = x.to(device) # y = y.to(device) # classifier = classifier.train() # pred = classifier(x) # predictions = torch.argmax(pred, dim=1) # total += y.size(0) # correct += (predictions == y).sum().item() # train_acc= train_correct / train_total # acc = correct / total # print(f"Training accuracy: {np.mean(train_acc)}") # print(f"Testing accuracy: {np.mean(acc)}") return train_instance_acc
def main(args): '''HYPER PARAMETER''' os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu '''CREATE DIR''' experiment_dir = Path('./experiment/') experiment_dir.mkdir(exist_ok=True) file_dir = Path(str(experiment_dir) + '/%s_ModelNet40-' % args.model_name + str( datetime.datetime.now().strftime('%Y-%m-%d_%H-%M'))) file_dir.mkdir(exist_ok=True) checkpoints_dir = file_dir.joinpath('checkpoints/') checkpoints_dir.mkdir(exist_ok=True) log_dir = file_dir.joinpath('logs/') log_dir.mkdir(exist_ok=True) '''LOG''' args = parse_args() logger = logging.getLogger(args.model_name) logger.setLevel(logging.INFO) formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s') file_handler = logging.FileHandler(str(log_dir) + 'train_%s_cls.txt' % args.model_name) file_handler.setLevel(logging.INFO) file_handler.setFormatter(formatter) logger.addHandler(file_handler) logger.info( '---------------------------------------------------TRANING---------------------------------------------------') logger.info('PARAMETER ...') logger.info(args) '''DATA LOADING''' logger.info('Load dataset ...') DATA_PATH = './data/modelnet40_normal_resampled/' TRAIN_DATASET = ModelNetDataLoader(root=DATA_PATH, npoint=args.num_point, split='train', normal_channel=args.normal) TEST_DATASET = ModelNetDataLoader(root=DATA_PATH, npoint=args.num_point, split='test', normal_channel=args.normal) trainDataLoader = torch.utils.data.DataLoader(TRAIN_DATASET, batch_size=args.batchsize, shuffle=True, num_workers=args.num_workers) testDataLoader = torch.utils.data.DataLoader(TEST_DATASET, batch_size=args.batchsize, shuffle=False, num_workers=args.num_workers) logger.info("The number of training data is: %d", len(TRAIN_DATASET)) logger.info("The number of test data is: %d", len(TEST_DATASET)) seed = 3 torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) '''MODEL LOADING''' num_class = 40 classifier = PointConvClsSsg(num_class).cuda() if args.pretrain is not None: print('Use pretrain model...') logger.info('Use pretrain model') checkpoint = torch.load(args.pretrain) start_epoch = checkpoint['epoch'] classifier.load_state_dict(checkpoint['model_state_dict']) else: print('No existing model, starting training from scratch...') start_epoch = 0 if args.optimizer == 'SGD': optimizer = torch.optim.SGD(classifier.parameters(), lr=0.01, momentum=0.9) elif args.optimizer == 'Adam': optimizer = torch.optim.Adam( classifier.parameters(), lr=args.learning_rate, betas=(0.9, 0.999), eps=1e-08, weight_decay=args.decay_rate ) scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=30, gamma=0.7) global_epoch = 0 global_step = 0 best_tst_accuracy = 0.0 blue = lambda x: '\033[94m' + x + '\033[0m' '''TRANING''' logger.info('Start training...') for epoch in range(start_epoch, args.epoch): print('Epoch %d (%d/%s):' % (global_epoch + 1, epoch + 1, args.epoch)) logger.info('Epoch %d (%d/%s):', global_epoch + 1, epoch + 1, args.epoch) mean_correct = [] scheduler.step() for batch_id, data in tqdm(enumerate(trainDataLoader, 0), total=len(trainDataLoader), smoothing=0.9): # for batch_id, data in enumerate(trainDataLoader, 0): points, target = data points = points.data.numpy() # 增强数据: 随机放大和平移点云,随机移除一些点 jittered_data = provider.random_scale_point_cloud(points[:, :, 0:3], scale_low=2.0 / 3, scale_high=3 / 2.0) jittered_data = provider.shift_point_cloud(jittered_data, shift_range=0.2) points[:, :, 0:3] = jittered_data points = provider.random_point_dropout_v2(points) provider.shuffle_points(points) points = torch.Tensor(points) target = target[:, 0] points = points.transpose(2, 1) points, target = points.cuda(), target.cuda() optimizer.zero_grad() classifier = classifier.train() # pred = classifier(points[:, :3, :], points[:, 3:, :]) pred = classifier(points[:, :3, :], None) loss = F.nll_loss(pred, target.long()) pred_choice = pred.data.max(1)[1] correct = pred_choice.eq(target.long().data).cpu().sum() mean_correct.append(correct.item() / float(points.size()[0])) loss.backward() optimizer.step() global_step += 1 train_acc = np.mean(mean_correct) print('Train Accuracy: %f' % train_acc) logger.info('Train Accuracy: %f' % train_acc) acc = test(classifier, testDataLoader) if (acc >= best_tst_accuracy) and epoch > 5: best_tst_accuracy = acc logger.info('Save model...') save_checkpoint( global_epoch + 1, train_acc, acc, classifier, optimizer, str(checkpoints_dir), args.model_name) print('Saving model....') print('\r Loss: %f' % loss.data) logger.info('Loss: %.2f', loss.data) print('\r Test %s: %f *** %s: %f' % (blue('Accuracy'), acc, blue('Best Accuracy'), best_tst_accuracy)) logger.info('Test Accuracy: %f *** Best Test Accuracy: %f', acc, best_tst_accuracy) global_epoch += 1 print('Best Accuracy: %f' % best_tst_accuracy) logger.info('End of training...')
def main(args): '''HYPER PARAMETER''' os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu '''CREATE DIR''' experiment_dir = Path('./eval_experiment/') experiment_dir.mkdir(exist_ok=True) file_dir = Path(str(experiment_dir) + '/%s_ModelNet40-'%args.model_name + str(datetime.datetime.now().strftime('%Y-%m-%d_%H-%M'))) file_dir.mkdir(exist_ok=True) checkpoints_dir = file_dir.joinpath('checkpoints/') checkpoints_dir.mkdir(exist_ok=True) os.system('cp %s %s' % (args.checkpoint, checkpoints_dir)) log_dir = file_dir.joinpath('logs/') log_dir.mkdir(exist_ok=True) '''LOG''' args = parse_args() logger = logging.getLogger(args.model_name) logger.setLevel(logging.INFO) formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s') file_handler = logging.FileHandler(str(log_dir) + 'eval_%s_cls.txt'%args.model_name) file_handler.setLevel(logging.INFO) file_handler.setFormatter(formatter) logger.addHandler(file_handler) logger.info('---------------------------------------------------EVAL---------------------------------------------------') logger.info('PARAMETER ...') logger.info(args) '''DATA LOADING''' logger.info('Load dataset ...') DATA_PATH = './data/modelnet40_normal_resampled/' TEST_DATASET = ModelNetDataLoader(root=DATA_PATH, npoint=args.num_point, split='test', normal_channel=args.normal) testDataLoader = torch.utils.data.DataLoader(TEST_DATASET, batch_size=args.batchsize, shuffle=False, num_workers=args.num_workers) logger.info("The number of test data is: %d", len(TEST_DATASET)) seed = 3 torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) '''MODEL LOADING''' num_class = 40 classifier = PointConvClsSsg(num_class).cuda() if args.checkpoint is not None: print('Load CheckPoint...') logger.info('Load CheckPoint') checkpoint = torch.load(args.checkpoint) start_epoch = checkpoint['epoch'] classifier.load_state_dict(checkpoint['model_state_dict']) else: print('Please load Checkpoint to eval...') sys.exit(0) start_epoch = 0 blue = lambda x: '\033[94m' + x + '\033[0m' '''EVAL''' logger.info('Start evaluating...') print('Start evaluating...') classifier = classifier.eval() mean_correct = [] for batch_id, data in tqdm(enumerate(testDataLoader, 0), total=len(testDataLoader), smoothing=0.9): pointcloud, target = data target = target[:, 0] points = pointcloud.permute(0, 2, 1) points, target = points.cuda(), target.cuda() with torch.no_grad(): pred = classifier(points[:, :3, :], points[:, 3:, :]) pred_choice = pred.data.max(1)[1] correct = pred_choice.eq(target.long().data).cpu().sum() mean_correct.append(correct.item()/float(points.size()[0])) accuracy = np.mean(mean_correct) print('Total Accuracy: %f'%accuracy) logger.info('Total Accuracy: %f'%accuracy) logger.info('End of evaluation...')
def main(args): def log_string(str): logger.info(str) print(str) '''HYPER PARAMETER''' os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu '''CREATE DIR''' timestr = str(datetime.datetime.now().strftime('%Y-%m-%d_%H-%M')) experiment_dir = Path('./log/') experiment_dir.mkdir(exist_ok=True) experiment_dir = experiment_dir.joinpath('classification') experiment_dir.mkdir(exist_ok=True) if args.log_dir is None: experiment_dir = experiment_dir.joinpath(timestr) else: experiment_dir = experiment_dir.joinpath(args.log_dir) experiment_dir.mkdir(exist_ok=True) checkpoints_dir = experiment_dir.joinpath('checkpoints/') checkpoints_dir.mkdir(exist_ok=True) log_dir = experiment_dir.joinpath('logs/') log_dir.mkdir(exist_ok=True) '''LOG''' args = parse_args() logger = logging.getLogger("Model") logger.setLevel(logging.INFO) formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s') file_handler = logging.FileHandler('%s/%s.txt' % (log_dir, args.model)) file_handler.setLevel(logging.INFO) file_handler.setFormatter(formatter) logger.addHandler(file_handler) log_string('PARAMETER ...') log_string(args) '''DATA LOADING''' log_string('Load dataset ...') DATA_PATH = 'data/modelnet40_normal_resampled/' TRAIN_DATASET = ModelNetDataLoader(root=DATA_PATH, npoint=args.num_point, split='train', normal_channel=args.normal) TEST_DATASET = ModelNetDataLoader(root=DATA_PATH, npoint=args.num_point, split='test', normal_channel=args.normal) trainDataLoader = torch.utils.data.DataLoader(TRAIN_DATASET, batch_size=args.batch_size, shuffle=True, num_workers=4) testDataLoader = torch.utils.data.DataLoader(TEST_DATASET, batch_size=args.batch_size, shuffle=False, num_workers=4) '''MODEL LOADING''' num_class = 40 MODEL = importlib.import_module(args.model) shutil.copy('./models/%s.py' % args.model, str(experiment_dir)) shutil.copy('./models/pointnet_util_psn.py', str(experiment_dir)) classifier = MODEL.get_model(num_class,normal_channel=args.normal).cuda() criterion = MODEL.get_loss().cuda() try: checkpoint = torch.load(str(experiment_dir) + '/checkpoints/best_model.pth') start_epoch = checkpoint['epoch'] classifier.load_state_dict(checkpoint['model_state_dict']) log_string('Use pretrain model') except: log_string('No existing model, starting training from scratch...') start_epoch = 0 if args.optimizer == 'Adam': optimizer = torch.optim.Adam( classifier.parameters(), lr=args.learning_rate, betas=(0.9, 0.999), eps=1e-08, weight_decay=args.decay_rate ) else: optimizer = torch.optim.SGD(classifier.parameters(), lr=0.01, momentum=0.9) scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=20, gamma=0.7) global_epoch = 0 global_step = 0 best_instance_acc = 0.0 best_class_acc = 0.0 mean_correct = [] best_epoch = 0 '''TRANING''' logger.info('Start training...') for epoch in range(start_epoch,args.epoch): log_string('Epoch %d (%d/%s):' % (global_epoch + 1, epoch + 1, args.epoch)) scheduler.step() for batch_id, data in tqdm(enumerate(trainDataLoader, 0), total=len(trainDataLoader), smoothing=0.9): points, target = data points = points.data.numpy() points = provider.random_point_dropout(points) points[:,:, 0:3] = provider.random_scale_point_cloud(points[:,:, 0:3]) points[:,:, 0:3] = provider.shift_point_cloud(points[:,:, 0:3]) points = torch.Tensor(points) target = target[:, 0] points = points.transpose(2, 1) points, target = points.cuda(), target.cuda() optimizer.zero_grad() classifier = classifier.train() pred, trans_feat = classifier(points, False) loss = criterion(pred, target.long(), trans_feat) pred_choice = pred.data.max(1)[1] correct = pred_choice.eq(target.long().data).cpu().sum() mean_correct.append(correct.item() / float(points.size()[0])) loss.backward() optimizer.step() global_step += 1 train_instance_acc = np.mean(mean_correct) log_string('Train Instance Accuracy: %f' % train_instance_acc) with torch.no_grad(): instance_acc, class_acc = test(classifier.eval(), testDataLoader) if (instance_acc >= best_instance_acc): best_instance_acc = instance_acc best_epoch = epoch + 1 if (class_acc >= best_class_acc): best_class_acc = class_acc log_string('Test Instance Accuracy: %f, Class Accuracy: %f'% (instance_acc, class_acc)) log_string('Best Instance Accuracy: %f, Class Accuracy: %f'% (best_instance_acc, best_class_acc)) if (instance_acc >= best_instance_acc): logger.info('Save model...') savepath = str(checkpoints_dir) + '/best_model.pth' log_string('Saving at %s'% savepath) state = { 'epoch': best_epoch, 'instance_acc': instance_acc, 'class_acc': class_acc, 'model_state_dict': classifier.state_dict(), 'optimizer_state_dict': optimizer.state_dict(), } torch.save(state, savepath) global_epoch += 1 logger.info('End of training...')
print('=====> Building new model...') torch.manual_seed(args.seed) print("Random Seed: ", args.seed) START_EPOCH = 0 acc_list = [0] print('Successfully built!') ######################################## ## Load data ######################################## print('======> Loading data') TRAIN_DATASET = ModelNetDataLoader(root=DATA_PATH, npoint=args.num_points, split='train', normal_channel=args.normal) TEST_DATASET = ModelNetDataLoader(root=DATA_PATH, npoint=args.num_points, split='test', normal_channel=args.normal) train_loader = torch.utils.data.DataLoader(TRAIN_DATASET, batch_size=args.batch_size, shuffle=True, num_workers=4, drop_last=True) test_loader = torch.utils.data.DataLoader(TEST_DATASET, batch_size=args.batch_size, shuffle=False, num_workers=4, drop_last=False)