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 vis(args): test_data, test_label = load_data(root, train=False) log.info(test_data=test_data.shape, test_label=test_label.shape) log.debug('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) model.cuda() log.info('Using multi GPU:', args.gpu) if args.pretrain is None: log.err('No pretrain model') return log.debug('Loading pretrain model...') checkpoint = torch.load(args.pretrain) model.load_state_dict(checkpoint) model.eval() log.info('Press space to exit, press Q for next frame') for idx in range(test_data.shape[0]): point_np = test_data[idx:idx + 1] gt = test_label[idx][0] points = torch.from_numpy(point_np) points = points.transpose(2, 1).cuda() pred, trans_feat = model(points) pred_choice = pred.data.max(1)[1] log.info(gt=class_names[gt], pred_choice=class_names[pred_choice.cpu().numpy().item()]) point_cloud = open3d.geometry.PointCloud() point_cloud.points = open3d.utility.Vector3dVector(point_np[0]) vis = open3d.visualization.VisualizerWithKeyCallback() vis.create_window() vis.get_render_option().background_color = np.asarray([0, 0, 0]) vis.add_geometry(point_cloud) vis.register_key_callback(32, lambda vis: exit()) vis.run() vis.destroy_window()
def main(args): ''' --- SELECT DEVICES --- ''' # Select either gpu or cpu device = torch.device("cuda" if args.cuda else "cpu") # Select among available GPUs if args.cuda: os.environ["CUDA_VISIBLE_DEVICES"] = ','.join( str(x) for x in args.gpudevice) ''' --- CREATE EXPERIMENTS DIRECTORY AND LOGGERS IN TENSORBOARD --- ''' projdir = sys.path[0] # Path for saving and loading the network. saveloadpath = os.path.join(projdir, 'experiment\\checkpoints', args.exp_name + '.pth') Path(os.path.dirname(saveloadpath)).mkdir(exist_ok=True, parents=True) # timestamp = str(datetime.datetime.now().strftime('%Y-%m-%d-%H-%M')) tblogdir = os.path.join(projdir, 'experiment\\tensorboardX', args.exp_name) # + '_' + timestamp ) Path(tblogdir).mkdir(exist_ok=True, parents=True) # Create tb_writer(the writer will be used to write the information on tb) by using SummaryWriter, # flush_secs defines how much seconds need to wait for writing information. tb_writer = SummaryWriter(logdir=tblogdir, flush_secs=3, write_to_disk=True) ''' --- INIT DATASETS AND DATALOADER (FOR SINGLE EPOCH) --- ''' # Read data from file, and create training data and testing data which are both in multiple frames. Beware Ts is # recording for every frame, i.e. every 82ms the automotive radar records once to form single frame(We need this information for LSTM). train_dataset, test_dataset, class_names = read_dataset( args.datapath, Ts=0.082, train_test_split=0.8) # Prepare the traing and testing dataset. both trainDataset and testDataset are dataset have multiple frames data, # for each frame it contains the "unified number of detection points"(NMAX detection points per frame). # Init test dataset(Beware we should NOT use data augmentation for test dataset) test_dataTransformations = transforms.Compose( [NormalizeTime(), Resampling(maxPointsPerFrame=10)]) testDataset = RadarClassDataset(dataset=test_dataset, transforms=test_dataTransformations, sequence_length=1) # Init train datasets train_dataTransformations = transforms.Compose([ NormalizeTime(), DataAugmentation(), Resampling(maxPointsPerFrame=10) ]) trainDataset = RadarClassDataset(dataset=train_dataset, transforms=train_dataTransformations, sequence_length=1) # Create dataloader for training by using batch_size frames' data in each batch trainDataLoader = DataLoader(trainDataset, batch_size=args.batchsize, shuffle=True, num_workers=args.num_workers) ''' --- INIT NETWORK MODEL --- ''' # Load selected network model and put it to right device if args.model_name == 'pointnet': classifier = PointNetCls(dim=args.pointCoordDim, num_class=args.numclasses, feature_transform=args.feature_transform) elif args.model_name == 'pointnet2': classifier = PointNet2ClsMsg( dim=args.pointCoordDim, num_class=args.numclasses, ) else: raise Exception( 'Argument "model_name" does not match existent networks') classifier = classifier.to(device) ''' --- LOAD NETWORK IF EXISTS --- ''' if os.path.exists(saveloadpath): print('Using pretrained model found...') checkpoint = torch.load(saveloadpath) start_epoch = checkpoint[ 'epoch'] + 1 # Just becase make sure counting starts from 1, 2, ..., rather than 0, 1, ..., when print the information of start_epoch iteration = checkpoint['iteration'] best_test_acc = checkpoint['test_accuracy'] classifier.load_state_dict(checkpoint['model_state_dict']) else: print('No existing model, starting training from scratch...') start_epoch = 1 # Just becase make sure counting starts from 1, 2, ..., rather than 0, 1, ..., when print the information of start_epoch iteration = 1 # Just becase make sure counting starts from 1, 2, ..., rather than 0, 1, ..., when print the information of iteration best_test_acc = 0 ''' --- CREATE OPTIMIZER ---''' if args.optimizer == 'SGD': optimizer = torch.optim.SGD(classifier.parameters(), lr=args.lr, momentum=0.9) elif args.optimizer == 'ADAM': optimizer = torch.optim.Adam(classifier.parameters(), lr=args.lr, betas=(0.9, 0.999), eps=1e-08, weight_decay=args.decay_rate) scheduler = torch.optim.lr_scheduler.StepLR( optimizer, step_size=args.lr_epoch_half, gamma=0.5) # half(0.5) the learning rate every 'step_size' epochs # log info printparams = 'Model parameters:' + json.dumps( vars(args), indent=4, sort_keys=True) print(printparams) tb_writer.add_text('hyper-parameters', printparams, iteration) # tb_writer.add_hparam(args) tb_writer.add_text( 'dataset', 'dataset sample size: training: {}, test: {}'.format( train_dataset.shape[0], test_dataset.shape[0]), iteration) ''' --- START TRANING ---''' for epoch in range(start_epoch, args.epoch + 1): print('Epoch %d/%s:' % (epoch, args.epoch)) # Add the "learning rate" into tensorboard scalar which will be shown in tensorboard tb_writer.add_scalar('learning_rate', optimizer.param_groups[0]['lr'], iteration) for batch_id, data in tqdm(enumerate(trainDataLoader, 0), total=len(trainDataLoader), smoothing=0.9): points, target = data # (B:batch x S:seq x C:features x N:points) , (B x S:seq) # Squeeze to drop Sequence dimension, which is equal to 1, convert all the data to float(otherwise there will be data type problems when running the model) and move to device points, target = points.squeeze( dim=1).float().to(device), target.float().to( device) # (B:batch x C:features x N:points) , (B) # points, target = points.float().to(device), target.float().to(device) # Reset gradients optimizer.zero_grad() # Sets the module in training mode classifier = classifier.train() # Forward propagation pred = classifier(points) # Calculate cross entropy loss (In the pointnet/pointnet2 network model, it outputs log_softmax result. Since # "log_softmax -> nll_loss" == CrossEntropyLoss, so that we just need to call F.nll_loss) loss = F.nll_loss(pred, target.long()) if args.model_name == 'pointnet': loss += feature_transform_regularizer(classifier.trans) * 0.001 if args.feature_transform: loss += feature_transform_regularizer( classifier.trans_feat) * 0.001 # Back propagate loss.backward() # Update weights optimizer.step() # Log once for every 5 batches, add the "train_loss/cross_entropy" into tensorboard scalar which will be shown in tensorboard if not batch_id % 5: tb_writer.add_scalar('train_loss/cross_entropy', loss.item(), iteration) iteration += 1 # if batch_id> 2: break scheduler.step() ''' --- TEST AND SAVE NETWORK --- ''' if not epoch % 10: # Doing the following things every epoch. # Perform predictions on the training data. train_targ, train_pred = test(classifier, trainDataset, device, num_workers=args.num_workers, batch_size=1800) # Perform predictions on the testing data. test_targ, test_pred = test(classifier, testDataset, device, num_workers=args.num_workers, batch_size=1800) # Calculate the accuracy rate for training data. train_acc = metrics_accuracy(train_targ, train_pred) # Calculate the accuracy rate for testing data. test_acc = metrics_accuracy(test_targ, test_pred) print('\r Training loss: {}'.format(loss.item())) print('Train Accuracy: {}\nTest Accuracy: {}'.format( train_acc, test_acc)) # Add the "train_acc" "test_acc" into tensorboard scalars which will be shown in tensorboard. tb_writer.add_scalars('metrics/accuracy', { 'train': train_acc, 'test': test_acc }, iteration) # Calculate confusion matrix. confmatrix_test = metrics_confusion_matrix(test_targ, test_pred) print('Test confusion matrix: \n', confmatrix_test) # Log confusion matrix. fig, ax = plot_confusion_matrix(confmatrix_test, class_names, normalize=False, title='Test Confusion Matrix') # Log normalized confusion matrix. fig_n, ax_n = plot_confusion_matrix( confmatrix_test, class_names, normalize=True, title='Test Confusion Matrix - Normalized') # Add the "confusion matrix" "normalized confusion matrix" into tensorboard figure which will be shown in tensorboard. tb_writer.add_figure('test_confusion_matrix/abs', fig, global_step=iteration, close=True) tb_writer.add_figure('test_confusion_matrix/norm', fig_n, global_step=iteration, close=True) # Log precision recall curves. for idx, clsname in enumerate(class_names): # Convert log_softmax to softmax(which is actual probability) and select the desired class. test_pred_binary = torch.exp(test_pred[:, idx]) test_targ_binary = test_targ.eq(idx) # Add the "precision recall curves" which will be shown in tensorboard. tb_writer.add_pr_curve(tag='pr_curves/' + clsname, labels=test_targ_binary, predictions=test_pred_binary, global_step=iteration) ''' --- SAVE NETWORK --- ''' # if (test_acc >= best_test_acc): # For now lets save every time, since we are only testing in a subset of the test dataset best_test_acc = test_acc # if test_acc > best_test_acc else best_test_acc state = { 'epoch': epoch, 'iteration': iteration, 'train_accuracy': train_acc if args.train_metric else 0.0, 'test_accuracy': best_test_acc, 'model_state_dict': classifier.state_dict(), 'optimizer_state_dict': optimizer.state_dict(), } torch.save(state, saveloadpath) print('Model saved!!!') print('Best Accuracy: %f' % best_test_acc) tb_writer.close()
def main(): ''' --- SELECT DEVICES --- ''' # Select either gpu or cpu device = torch.device("cuda" if args.cuda else "cpu") # Select among available GPUs if args.cuda: os.environ["CUDA_VISIBLE_DEVICES"] = ','.join( str(x) for x in args.gpudevice) ''' --- INIT NETWORK MODEL --- ''' projdir = sys.path[0] # Path for saving and loading the network. saveloadpath = os.path.join(projdir, 'experiment\\checkpoints', args.exp_name + '.pth') Path(os.path.dirname(saveloadpath)).mkdir(exist_ok=True, parents=True) # Load selected network model and put it to right device if args.model_name == 'pointnet': classifier = PointNetCls(dim=args.pointCoordDim, num_class=len(args.categories), feature_transform=args.feature_transform) elif args.model_name == 'pointnet2': classifier = PointNet2ClsMsg(dim=args.pointCoordDim, num_class=len(args.categories)) else: raise Exception( 'Argument "model_name" does not match existent networks') classifier = classifier.to(device) ''' --- LOAD NETWORK IF EXISTS --- ''' if os.path.exists(saveloadpath): print('Using pretrained model found...') checkpoint = torch.load(saveloadpath) start_epoch = checkpoint[ 'epoch'] + 1 # Just becase make sure counting starts from 1, 2, ..., rather than 0, 1, ..., when print the information of start_epoch iteration = checkpoint['iteration'] best_test_acc = checkpoint['test_accuracy'] classifier.load_state_dict(checkpoint['model_state_dict']) else: raise Exception('Model in: {} does not exists'.format(saveloadpath)) # put classifier in evaluation mode classifier = classifier.eval() ''' --- INIT DATASETS AND DATALOADER (FOR SINGLE EPOCH) --- ''' # Ideal for PointNet and pointLSTM - dataloader will return (B:batch, S:seq, C:features, N:points) dataTransformations = transforms.Compose([ ToSeries(), DataAugmentation(), Resampling(maxPointsPerFrame=10), ToTensor() ]) # init dataset nusc = NuScenes(version=args.nuscenes_eval_dir, dataroot=args.nuscenes_dir, verbose=True) ''' Iterate over samples ''' idx = list(range(len(nusc.sample))) random.shuffle(idx) # shuffle samples # for sample_rec in nusc.sample: for i in idx: sample_rec = nusc.sample[i] sample_token = sample_rec['token'] # read all sensors data and merge them to a common reference frame sensorlist = [] for sensor in args.sensors: sensorlist.append( MyRadarPointCloud.from_sample_token(nusc, sample_token, sensor)) global radar_pc radar_pc = MyRadarPointCloud.merge(sensorlist, 0) # get annotations for actual sample ann_tokens = nusc.get('sample', sample_token)['anns'] # filter annotations that are objects of interest for our classifier def get_ann_properties(nusc, ann_token: str): ann_rec = nusc.get('sample_annotation', ann_token) category = ann_rec['category_name'] attr_tokens = ann_rec['attribute_tokens'] attr = [ nusc.get('attribute', attr_token)['name'] for attr_token in attr_tokens ] if len(attr) == 0: attr = [''] print('object category:{:30} attr:{:40} center:{:20}'.format( category, str(attr), str(ann_rec['translation'] - radar_pc.A_cs_2_gl[:3, 3]))) return category, attr[ 0] # for now lets return only the first attribute (we might have more per object) def isInterestingObject(ann_token): cat, att = get_ann_properties(nusc, ann_token) # category-attribute pair must match at least one of the desired category-attr pair for desired_cat, desired_att in zip(args.categories, args.attributes): if cat in desired_cat and att in desired_att: return True return False print('\n\nGround-truth objects on this frame:') ann_tokens = list(filter(isInterestingObject, ann_tokens)) def getObjLabel(ann_token): cat, att = get_ann_properties(nusc, ann_token) for idx, (desired_cat, desired_att) in enumerate( zip(args.categories, args.attributes)): if cat in desired_cat and att in desired_att: return idx return np.NaN labels = list(map(getObjLabel, ann_tokens)) assert np.all( ~np.isnan(labels) ), 'Something strange happened... object was selected as interesting but we can not find its label' # create bounding boxes from annotations ann_boxes = [] for ann_token, label in zip(ann_tokens, labels): box = radar_pc.box(ann_token) box.label = label box.name = args.objectNames[label] ann_boxes.append(box) ''' APPLY DBSCAN ON EACH SCENE''' from sklearn.cluster import DBSCAN points_scene = radar_pc.points # points_gl = <5,N> # filter out objects, whose speed is less than 0.3 m/s idx_moving = np.linalg.norm(points_scene[2:4, :].T, axis=1) > 0.25 points_scene = points_scene[:, idx_moving] # apply DBSCAN clustering = DBSCAN(eps=3, min_samples=1).fit(points_scene[:2, :].T) ''' For each cluster, run network and predict class ''' pred_results = [] for cluster_idx in range(max(clustering.labels_)): # select points from the current cluster points_idx = clustering.labels_ == cluster_idx points_obj = points_scene[:, points_idx] # apply necessary transformations features = dataTransformations(points_obj) # convert to torch.tensor features = torch.tensor(features).type( torch.FloatTensor).unsqueeze(0).to(device) # calculate network prediction pred = classifier(features).argmax().item() # store result pred_results.append((points_obj, pred)) ''' RENDER ''' OBJCOLORS = ['red', 'green', 'blue', 'grey', 'orange', 'white'] fig = plt.figure(constrained_layout=True) gs = GridSpec(2, 6, figure=fig) gs.update(wspace=0, hspace=0) axs = [ fig.add_subplot(gs[1, :3]), fig.add_subplot(gs[1, 3:]), fig.add_subplot(gs[0, 0]), fig.add_subplot(gs[0, 1]), fig.add_subplot(gs[0, 2]), fig.add_subplot(gs[0, 3]), fig.add_subplot(gs[0, 4]), fig.add_subplot(gs[0, 5]) ] # render annotation boxes for box in ann_boxes: color = OBJCOLORS[box.label] box.render(axs[0], colors=(color, ) * 3) box.render(axs[1], colors=(color, ) * 3) # render cars radar_pc.car.render(axs[0], colors=('orange', 'k', 'k')) radar_pc.car.render(axs[1], colors=('orange', 'k', 'k')) # render pointcloud radar_pc._render_pc(axs[0], radar_pc.points, color_channel='k') for cluster_points, pred in pred_results: # radar_pc._render_pc( axs[1], cluster_points, color_channel=OBJCOLORS[pred]) if pred != 5: axs[1].scatter(cluster_points[0, :], cluster_points[1, :], s=30, c=OBJCOLORS[pred], edgecolors='k', linewidths=1, zorder=100) # plot ellipses if cluster_points.shape[1] > 1: confidence_ellipse(axs[1], cluster_points[0, :], cluster_points[1, :], edgecolor=OBJCOLORS[pred], facecolor=OBJCOLORS[pred], alpha=0.5) # render camera images cameras = [ 'CAM_BACK_LEFT', 'CAM_FRONT_LEFT', 'CAM_FRONT', 'CAM_FRONT_RIGHT', 'CAM_BACK_RIGHT', 'CAM_BACK' ] for ax_idx, cam_name in enumerate(cameras): cam_token = sample_rec['data'][cam_name] cam_rec = nusc.get('sample_data', cam_token) img_filename = os.path.join(nusc.dataroot, cam_rec['filename']) img = mpimg.imread(img_filename) # axs[ax_idx].set_axis_off() axs[ax_idx + 2].get_xaxis().set_visible(False) axs[ax_idx + 2].get_yaxis().set_visible(False) axs[ax_idx + 2].imshow(img) axs[ax_idx + 2].set_title(cam_name) # format legend_elements = [ Line2D([0], [0], color=color, lw=4, label=name) for color, name in zip(OBJCOLORS, args.objectNames) ] axs[0].legend(handles=legend_elements, loc='upper right', prop={'size': 8}) axs[1].legend(handles=legend_elements, loc='upper right', prop={'size': 8}) axs[0].axis('equal') axs[1].axis('equal') mng = plt.get_current_fig_manager() mng.window.showMaximized() plt.tight_layout() lims = (-60, 60) axs[0].set_xlim(lims) axs[0].set_ylim(lims) axs[0].set_title( 'Radar detection points and ground truth bounding boxes of MOVING objects' ) axs[1].set_xlim(lims) axs[1].set_ylim(lims) axs[1].set_title( 'Segmented segmentation of detection points (DBSCAN + PointNet classifier)\n3-sigma ellipses for each cluster and ground-truth bbs' ) plt.tight_layout() plt.show() # plot scene with the standard Nuscenes method # nusc.render_sample(sample_token) # plt.show() # sample_token = nusc.get('sample', sample_token)['next'] tb_writer.close()
num_workers=16, pin_memory=True, shuffle=True, worker_init_fn=worker_init_fn, collate_fn=my_collate, ) is_resume = 0 if args.load_model and args.load_epoch != -1: is_resume = 1 model = PointNetCls(num_points=grasp_points_num, input_chann=point_channel, k=2) if is_resume or args.mode == 'test': model.load_state_dict(torch.load(args.load_model)) print('load weights {}'.format(args.load_model)) if args.cuda: base_model = model.cuda() model = nn.DataParallel(base_model) optimizer = optim.Adam(model.parameters(), lr=args.lr) scheduler = StepLR(optimizer, step_size=30, gamma=0.5) def train(model, dataset, loader, epoch): scheduler.step() dataset.train() model.train() torch.set_grad_enabled(True) correct = 0 dataset_size = 0
def main(args): '''HYPER PARAMETER''' os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu # datapath = './data/ModelNet/' datapath = './data/objecnn20_data_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) file_dir = Path(str(experiment_dir) +'/%sObjectNNClf-'%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_ObjectNNClf.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 = ObjectNNDataLoader(train_data, train_label, rotation=ROTATION) if ROTATION is not None: print('The range of training rotation is',ROTATION) testDataset = ObjectNNDataLoader(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 = 20 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']) # print(checkpoint['model_state_dict']) model_dict = classifier.state_dict() pretrained_dict = {k: v for k, v in checkpoint['model_state_dict'].items() if k in model_dict} model_dict.update(pretrained_dict) classifier.load_state_dict(model_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=20, 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, trans_feat, global_feature = classifier(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 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' % (blue('Accuracy'),acc)) logger.info('Test Accuracy: %f', acc) 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(testDataset): args = parse_args() '''HYPER PARAMETER''' # os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu if args.multi_gpu is None else '0,1,2,3' datapath = './data/ModelNet/' 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/train_%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('---------------------------------------------------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, 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 classifier = PointNetCls(num_class,args.feature_transform).cuda() if args.model_name == 'pointnet' else PointNet2ClsMsg().cuda() '''GPU selection and multi-GPU''' if args.multi_gpu is not None: device_ids = [int(x) for x in args.multi_gpu.split(',')] torch.backends.cudnn.benchmark = True classifier.cuda(device_ids[0]) classifier = torch.nn.DataParallel(classifier, device_ids=device_ids) else: classifier.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=20, gamma=0.5) global_epoch = 0 global_step = 0 best_tst_accuracy = 0.0 blue = lambda x: '\033[94m' + x + '\033[0m' start_epoch = 0 args.epoch = 1 '''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, trans_feat = classifier(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 train_acc = test(classifier.eval(), trainDataLoader) if args.train_metric else None acc, fts = test(classifier, testDataLoader) # return fts # 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' % (blue('Accuracy'),acc)) logger.info('Test Accuracy: %f', acc) return fts 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 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(): ''' --- SELECT DEVICES --- ''' # Select either gpu or cpu device = torch.device("cuda" if args.cuda else "cpu") # Select among available GPUs if args.cuda: os.environ["CUDA_VISIBLE_DEVICES"] = ','.join(str(x) for x in args.gpudevice) ''' --- CREATE EXPERIMENTS DIRECTORY AND LOGGERS IN TENSORBOARD --- ''' projdir = sys.path[0] # Path for saving and loading the network. saveloadpath = os.path.join( projdir, 'experiment\\checkpoints', args.exp_name+'.pth') Path(os.path.dirname(saveloadpath)).mkdir(exist_ok=True, parents=True) # timestamp = str(datetime.datetime.now().strftime('%Y-%m-%d-%H-%M')) tblogdir = os.path.join( projdir, 'experiment\\tensorboardX', args.exp_name ) # + '_' + timestamp ) Path(tblogdir).mkdir(exist_ok=True, parents=True) # Create tb_writer(the writer will be used to write the information on tb) by using SummaryWriter, # flush_secs defines how much seconds need to wait for writing information. tb_writer = SummaryWriter( logdir=tblogdir, flush_secs=3, write_to_disk=True) ''' --- INIT DATASETS AND DATALOADER (FOR SINGLE EPOCH) --- ''' # Ideal for PointNet and pointLSTM - dataloader will return (B:batch, S:seq, C:features, N:points) dataTransformations = transforms.Compose([ ToSeries(), DataAugmentation(), Resampling(maxPointsPerFrame=10), ToTensor() ]) # Init nuScenes datasets nusc_train = NuScenes(version=args.nuscenes_train_dir, dataroot=args.nuscenes_dir, verbose=True) train_dataset = RadarClassDataset(nusc_train, categories=args.categories, sensors=args.sensors, transforms=dataTransformations, sequence_length=1) nusc_test = NuScenes(version=args.nuscenes_test_dir, dataroot=args.nuscenes_dir, verbose=True) test_dataset = RadarClassDataset(nusc_test, categories=args.categories, sensors=args.sensors, transforms=dataTransformations, sequence_length=1) # Init training data loader trainDataLoader = DataLoader(train_dataset, batch_size=args.batchsize, shuffle=True, num_workers=args.num_workers) ''' --- INIT NETWORK MODEL --- ''' # Load selected network model and put it to right device if args.model_name == 'pointnet': classifier = PointNetCls(dim=args.pointCoordDim, num_class=len(args.categories), feature_transform=args.feature_transform) elif args.model_name == 'pointnet2': classifier = PointNet2ClsMsg(dim=args.pointCoordDim, num_class=len(args.categories) ) else: raise Exception('Argument "model_name" does not match existent networks') classifier = classifier.to(device) ''' --- INIT LOSS FUNCTION --- ''' loss_fun = FocalLoss(gamma=args.focalLoss_gamma, num_classes=len(args.categories), alpha=args.weight_cat).to(device) ''' --- LOAD NETWORK IF EXISTS --- ''' if os.path.exists(saveloadpath): print('Using pretrained model found...') checkpoint = torch.load(saveloadpath) start_epoch = checkpoint['epoch'] +1 # Just becase make sure counting starts from 1, 2, ..., rather than 0, 1, ..., when print the information of start_epoch iteration = checkpoint['iteration'] best_test_acc = checkpoint['test_accuracy'] classifier.load_state_dict(checkpoint['model_state_dict']) else: print('No existing model, starting training from scratch...') start_epoch = 1 # Just becase make sure counting starts from 1, 2, ..., rather than 0, 1, ..., when print the information of start_epoch iteration = 1 # Just becase make sure counting starts from 1, 2, ..., rather than 0, 1, ..., when print the information of iteration best_test_acc = 0 ''' --- CREATE OPTIMIZER ---''' if args.optimizer == 'SGD': optimizer = torch.optim.SGD( classifier.parameters(), lr=args.lr, momentum=0.9) elif args.optimizer == 'ADAM': optimizer = torch.optim.Adam( classifier.parameters(), lr=args.lr, betas=(0.9, 0.999), eps=1e-08, weight_decay=args.decay_rate) scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=args.lr_epoch_half, gamma=0.5) # half(0.5) the learning rate every 'step_size' epochs # Log info printparams = 'Model parameters:' + json.dumps(vars(args), indent=4, sort_keys=True) print(printparams) tb_writer.add_text('hyper-parameters',printparams,iteration) # tb_writer.add_hparam(args) tb_writer.add_text('dataset','dataset sample size: training: {}, test: {}'.format(len(train_dataset),len(test_dataset)),iteration) ''' --- START TRANING ---''' for epoch in range(start_epoch, args.epoch+1): # epoch = start_epoch print('Epoch %d/%s:' % (epoch, args.epoch)) # Add the "learning rate" into tensorboard scalar which will be shown in tensorboard tb_writer.add_scalar('learning_rate', optimizer.param_groups[0]['lr'], iteration) # Beware epochs_left = args.epoch - epoch for batch_id, data in tqdm(enumerate(trainDataLoader, 0), total=len(trainDataLoader), smoothing=0.9): points, target = data # (B:batch x S:seq x C:features x N:points) , (B x S:seq) # Squeeze to drop Sequence dimension, which is equal to 1, convert all the data to float(otherwise there will be data type problems when running the model) and move to device points, target = points.squeeze(dim=1).float().to(device), target.float().to(device) # (B:batch x C:features x N:points) , (B) # points, target = points.float().to(device), target.float().to(device) # Reset gradients optimizer.zero_grad() # Sets the module in training mode classifier = classifier.train() # Forward propagation pred = classifier(points) # MLE estimator = min (- log (softmax(x)) ) = min nll_loss(log_softmax(x)) # loss = F.nll_loss(pred, target.long()) loss = loss_fun(pred, target.long()) if args.model_name == 'pointnet': loss += feature_transform_regularizer(classifier.trans) * 0.001 if args.feature_transform: loss += feature_transform_regularizer(classifier.trans_feat) * 0.001 # Back propagate loss.backward() # Update weights optimizer.step() # Log once for every 5 batches, add the "train_loss/cross_entropy" into tensorboard scalar which will be shown in tensorboard if not batch_id % 5: tb_writer.add_scalar('train_loss/cross_entropy', loss.item(), iteration) iteration += 1 # Plot train confusion matrix every X steps if not iteration % 20: confmatrix_train = metrics_confusion_matrix(target, pred) print('\nTrain confusion matrix: \n',confmatrix_train) # We just finished one epoch # if not batch_id+1 % int(train_dataset.len__()/args.batchsize): ''' --- TEST NETWORK --- ''' if not epoch % int(args.test_every_X_epochs): # Doing the following things every epoch. # Perform predictions on the training data. train_targ, train_pred = test(classifier, train_dataset, device, num_workers=0, batch_size=512) # Perform predictions on the testing data. test_targ, test_pred = test(classifier, test_dataset, device, num_workers=0, batch_size=512) # Calculate the accuracy rate for training data. train_acc = metrics_accuracy(train_targ, train_pred) # Calculate the accuracy rate for testing data. test_acc = metrics_accuracy(test_targ, test_pred) print('\r Training loss: {}'.format(loss.item())) print('Train Accuracy: {}\nTest Accuracy: {}'.format(train_acc, test_acc) ) # Add the "train_acc" "test_acc" into tensorboard scalars which will be shown in tensorboard. tb_writer.add_scalars('metrics/accuracy', {'train':train_acc, 'test':test_acc}, iteration) # Calculate confusion matrix. confmatrix_test = metrics_confusion_matrix(test_targ, test_pred) print('Test confusion matrix: \n',confmatrix_test) # Log confusion matrix. fig, ax = plot_confusion_matrix(confmatrix_test, args.categories, normalize=False, title='Test Confusion Matrix') # Log normalized confusion matrix. fig_n, ax_n = plot_confusion_matrix(confmatrix_test, args.categories, normalize=True, title='Test Confusion Matrix - Normalized') # Add the "confusion matrix" "normalized confusion matrix" into tensorboard figure which will be shown in tensorboard. tb_writer.add_figure('test_confusion_matrix/abs', fig, global_step=iteration, close=True) tb_writer.add_figure('test_confusion_matrix/norm', fig_n, global_step=iteration, close=True) # Log precision recall curves. for idx, clsname in enumerate(args.categories): # Convert log_softmax to softmax(which is actual probability) and select the desired class. test_pred_binary = torch.exp(test_pred[:,idx]) test_targ_binary = test_targ.eq(idx) # Add the "precision recall curves" which will be shown in tensorboard. tb_writer.add_pr_curve(tag='pr_curves/'+clsname, labels=test_targ_binary, predictions=test_pred_binary, global_step=iteration) # Store the best test accuracy if (test_acc >= best_test_acc): best_test_acc = max([best_test_acc, test_acc]) # NOTE: we possibly want to save only when when the best test accuracy is surpassed. For now lets save every X epoch ''' --- SAVE NETWORK --- ''' if not epoch % int(args.save_every_X_epochs): print('Best Accuracy: %f'%best_test_acc) state = { 'epoch': epoch, 'iteration': iteration, 'test_accuracy': best_test_acc, 'model_state_dict': classifier.state_dict(), 'optimizer_state_dict': optimizer.state_dict(), } torch.save(state, saveloadpath) print('Model saved!!!') # epoch += 1 # print('Epoch %d/%s:' % (epoch, args.epoch)) scheduler.step() tb_writer.close()
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): '''HYPER PARAMETER''' os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu # datapath = './data/ModelNet/' datapath = './data/objecnn20_data_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) file_dir = Path( str(experiment_dir) + '/Test_%sObjectNNClf-' % 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) + '/test_%s_ObjectNNClf.txt' % args.model_name) 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 ...') database_data, database_label, query_data, query_label = load_data( datapath, classification=True) print(">>>>>>>>>database_data:", database_data.shape) print(">>>>>>>>>query_data:", query_data.shape) logger.info("The number of database_data data is: %d", database_data.shape[0]) logger.info("The number of query_data data is: %d", query_data.shape[0]) ###################### 加载 database 和 query ###################### databaseDataset = TestQueryObjectNNDataLoader(database_data, database_label, rotation=ROTATION) if ROTATION is not None: print('The range of training rotation is', ROTATION) queryDataset = TestQueryObjectNNDataLoader(query_data, query_label, rotation=ROTATION) databaseDataLoader = torch.utils.data.DataLoader(databaseDataset, batch_size=args.batchsize, shuffle=False) queryDataLoader = torch.utils.data.DataLoader(queryDataset, batch_size=args.batchsize, shuffle=False) # 不打乱 '''MODEL LOADING''' num_class = 20 ###################### PointNetCls ###################### classifier = PointNetCls(num_class, args.feature_transform).cuda( ) if args.model_name == 'pointnet' else PointNet2ClsMsg().cuda() # classifier = PointNetCls(num_class,args.feature_transform) if args.model_name == 'pointnet' else PointNet2ClsMsg() 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' '''QueryING''' logger.info('Start query...') scheduler.step() ###################### Query ###################### classifier.eval() _acc, database_feature_martix = getGlobalFeature('database', classifier.eval(), databaseDataLoader) _acc, query_feature_martix = getGlobalFeature('query', classifier.eval(), queryDataLoader)