From f680d6f7a880ff87bffa50d3d69a963a6298b541 Mon Sep 17 00:00:00 2001 From: Laurent FAINSIN Date: Thu, 3 Aug 2023 16:40:14 +0200 Subject: [PATCH] auto-formatting --- analysis.py | 11 +- classification_ModelNet40/data.py | 76 ++-- classification_ModelNet40/helper.py | 6 +- classification_ModelNet40/main.py | 197 +++++---- classification_ModelNet40/models/__init__.py | 2 - classification_ModelNet40/models/pointmlp.py | 280 ++++++++----- classification_ModelNet40/test.py | 79 ++-- classification_ModelNet40/utils/__init__.py | 5 +- classification_ModelNet40/utils/logger.py | 86 ++-- classification_ModelNet40/utils/misc.py | 152 +++---- .../utils/progress/progress/__init__.py | 14 +- .../utils/progress/progress/bar.py | 51 ++- .../utils/progress/progress/counter.py | 9 +- .../utils/progress/progress/helpers.py | 42 +- .../utils/progress/progress/spinner.py | 16 +- .../utils/progress/setup.py | 40 +- .../utils/progress/test_progress.py | 26 +- classification_ModelNet40/voting.py | 132 +++--- classification_ScanObjectNN/ScanObjectNN.py | 40 +- classification_ScanObjectNN/main.py | 199 +++++---- .../models/__init__.py | 2 - .../models/pointmlp.py | 281 ++++++++----- classification_ScanObjectNN/utils/__init__.py | 5 +- classification_ScanObjectNN/utils/logger.py | 86 ++-- classification_ScanObjectNN/utils/misc.py | 152 +++---- .../utils/progress/progress/__init__.py | 14 +- .../utils/progress/progress/bar.py | 51 ++- .../utils/progress/progress/counter.py | 9 +- .../utils/progress/progress/helpers.py | 42 +- .../utils/progress/progress/spinner.py | 16 +- .../utils/progress/setup.py | 40 +- .../utils/progress/test_progress.py | 26 +- part_segmentation/main.py | 381 +++++++++++------- part_segmentation/model/__init__.py | 1 - part_segmentation/model/pointMLP.py | 336 +++++++++------ part_segmentation/util/data_util.py | 95 +++-- part_segmentation/util/util.py | 35 +- .../pointnet2_ops/pointnet2_modules.py | 68 ++-- .../pointnet2_ops/pointnet2_utils.py | 100 ++--- pointnet2_ops_lib/setup.py | 4 +- 40 files changed, 1822 insertions(+), 1385 deletions(-) diff --git a/analysis.py b/analysis.py index 9eb1022..46fdc60 100644 --- a/analysis.py +++ b/analysis.py @@ -1,20 +1,21 @@ -import torch -import fvcore.nn import fvcore.common +import fvcore.nn +import torch from fvcore.nn import FlopCountAnalysis + from classification_ScanObjectNN.models import pointMLPElite model = pointMLPElite() model.eval() # model = deit_tiny_patch16_224() -inputs = (torch.randn((1,3,1024))) +inputs = torch.randn((1, 3, 1024)) k = 1024.0 flops = FlopCountAnalysis(model, inputs).total() print(f"Flops : {flops}") -flops = flops/(k**3) +flops = flops / (k**3) print(f"Flops : {flops:.1f}G") params = fvcore.nn.parameter_count(model)[""] print(f"Params : {params}") -params = params/(k**2) +params = params / (k**2) print(f"Params : {params:.1f}M") diff --git a/classification_ModelNet40/data.py b/classification_ModelNet40/data.py index ee3c5c2..7139372 100644 --- a/classification_ModelNet40/data.py +++ b/classification_ModelNet40/data.py @@ -1,33 +1,37 @@ -import os import glob +import os + import h5py import numpy as np from torch.utils.data import Dataset + os.environ["HDF5_USE_FILE_LOCKING"] = "FALSE" + def download(): BASE_DIR = os.path.dirname(os.path.abspath(__file__)) - DATA_DIR = os.path.join(BASE_DIR, 'data') + DATA_DIR = os.path.join(BASE_DIR, "data") if not os.path.exists(DATA_DIR): os.mkdir(DATA_DIR) - if not os.path.exists(os.path.join(DATA_DIR, 'modelnet40_ply_hdf5_2048')): - www = 'https://shapenet.cs.stanford.edu/media/modelnet40_ply_hdf5_2048.zip' + if not os.path.exists(os.path.join(DATA_DIR, "modelnet40_ply_hdf5_2048")): + www = "https://shapenet.cs.stanford.edu/media/modelnet40_ply_hdf5_2048.zip" zipfile = os.path.basename(www) - os.system('wget %s --no-check-certificate; unzip %s' % (www, zipfile)) - os.system('mv %s %s' % (zipfile[:-4], DATA_DIR)) - os.system('rm %s' % (zipfile)) + os.system(f"wget {www} --no-check-certificate; unzip {zipfile}") + os.system(f"mv {zipfile[:-4]} {DATA_DIR}") + os.system("rm %s" % (zipfile)) + def load_data(partition): download() BASE_DIR = os.path.dirname(os.path.abspath(__file__)) - DATA_DIR = os.path.join(BASE_DIR, 'data') + DATA_DIR = os.path.join(BASE_DIR, "data") all_data = [] all_label = [] - for h5_name in glob.glob(os.path.join(DATA_DIR, 'modelnet40_ply_hdf5_2048', 'ply_data_%s*.h5'%partition)): + for h5_name in glob.glob(os.path.join(DATA_DIR, "modelnet40_ply_hdf5_2048", "ply_data_%s*.h5" % partition)): # print(f"h5_name: {h5_name}") - f = h5py.File(h5_name,'r') - data = f['data'][:].astype('float32') - label = f['label'][:].astype('int64') + f = h5py.File(h5_name, "r") + data = f["data"][:].astype("float32") + label = f["label"][:].astype("int64") f.close() all_data.append(data) all_label.append(label) @@ -35,40 +39,42 @@ def load_data(partition): all_label = np.concatenate(all_label, axis=0) return all_data, all_label + def random_point_dropout(pc, max_dropout_ratio=0.875): - ''' batch_pc: BxNx3 ''' + """batch_pc: BxNx3.""" # for b in range(batch_pc.shape[0]): - dropout_ratio = np.random.random()*max_dropout_ratio # 0~0.875 - drop_idx = np.where(np.random.random((pc.shape[0]))<=dropout_ratio)[0] + dropout_ratio = np.random.random() * max_dropout_ratio # 0~0.875 + drop_idx = np.where(np.random.random(pc.shape[0]) <= dropout_ratio)[0] # print ('use random drop', len(drop_idx)) - if len(drop_idx)>0: - pc[drop_idx,:] = pc[0,:] # set to the first point + if len(drop_idx) > 0: + pc[drop_idx, :] = pc[0, :] # set to the first point return pc + def translate_pointcloud(pointcloud): - xyz1 = np.random.uniform(low=2./3., high=3./2., size=[3]) + xyz1 = np.random.uniform(low=2.0 / 3.0, high=3.0 / 2.0, size=[3]) xyz2 = np.random.uniform(low=-0.2, high=0.2, size=[3]) - - translated_pointcloud = np.add(np.multiply(pointcloud, xyz1), xyz2).astype('float32') - return translated_pointcloud + + return np.add(np.multiply(pointcloud, xyz1), xyz2).astype("float32") + def jitter_pointcloud(pointcloud, sigma=0.01, clip=0.02): N, C = pointcloud.shape - pointcloud += np.clip(sigma * np.random.randn(N, C), -1*clip, clip) + pointcloud += np.clip(sigma * np.random.randn(N, C), -1 * clip, clip) return pointcloud class ModelNet40(Dataset): - def __init__(self, num_points, partition='train'): + def __init__(self, num_points, partition="train"): self.data, self.label = load_data(partition) self.num_points = num_points - self.partition = partition + self.partition = partition def __getitem__(self, item): - pointcloud = self.data[item][:self.num_points] + pointcloud = self.data[item][: self.num_points] label = self.label[item] - if self.partition == 'train': + if self.partition == "train": # pointcloud = random_point_dropout(pointcloud) # open for dgcnn not for our idea for all pointcloud = translate_pointcloud(pointcloud) np.random.shuffle(pointcloud) @@ -78,19 +84,25 @@ class ModelNet40(Dataset): return self.data.shape[0] -if __name__ == '__main__': +if __name__ == "__main__": train = ModelNet40(1024) - test = ModelNet40(1024, 'test') + test = ModelNet40(1024, "test") # for data, label in train: # print(data.shape) # print(label.shape) from torch.utils.data import DataLoader - train_loader = DataLoader(ModelNet40(partition='train', num_points=1024), num_workers=4, - batch_size=32, shuffle=True, drop_last=True) + + train_loader = DataLoader( + ModelNet40(partition="train", num_points=1024), + num_workers=4, + batch_size=32, + shuffle=True, + drop_last=True, + ) for batch_idx, (data, label) in enumerate(train_loader): print(f"batch_idx: {batch_idx} | data shape: {data.shape} | ;lable shape: {label.shape}") - train_set = ModelNet40(partition='train', num_points=1024) - test_set = ModelNet40(partition='test', num_points=1024) + train_set = ModelNet40(partition="train", num_points=1024) + test_set = ModelNet40(partition="test", num_points=1024) print(f"train_set size {train_set.__len__()}") print(f"test_set size {test_set.__len__()}") diff --git a/classification_ModelNet40/helper.py b/classification_ModelNet40/helper.py index 5a31f89..cd51a1e 100644 --- a/classification_ModelNet40/helper.py +++ b/classification_ModelNet40/helper.py @@ -1,9 +1,9 @@ import torch import torch.nn.functional as F -def cal_loss(pred, gold, smoothing=True): - ''' Calculate cross entropy loss, apply label smoothing if needed. ''' +def cal_loss(pred, gold, smoothing=True): + """Calculate cross entropy loss, apply label smoothing if needed.""" gold = gold.contiguous().view(-1) if smoothing: @@ -16,6 +16,6 @@ def cal_loss(pred, gold, smoothing=True): loss = -(one_hot * log_prb).sum(dim=1).mean() else: - loss = F.cross_entropy(pred, gold, reduction='mean') + loss = F.cross_entropy(pred, gold, reduction="mean") return loss diff --git a/classification_ModelNet40/main.py b/classification_ModelNet40/main.py index 2d0ef92..d0c750d 100644 --- a/classification_ModelNet40/main.py +++ b/classification_ModelNet40/main.py @@ -1,41 +1,46 @@ -""" -Usage: -python main.py --model PointMLP --msg demo +"""Usage: +python main.py --model PointMLP --msg demo. """ import argparse -import os -import logging import datetime +import logging +import os + +import models as models +import numpy as np +import sklearn.metrics as metrics import torch -import torch.nn.parallel import torch.backends.cudnn as cudnn +import torch.nn.parallel import torch.optim import torch.utils.data import torch.utils.data.distributed -from torch.utils.data import DataLoader -import models as models -from utils import Logger, mkdir_p, progress_bar, save_model, save_args, cal_loss from data import ModelNet40 from torch.optim.lr_scheduler import CosineAnnealingLR -import sklearn.metrics as metrics -import numpy as np +from torch.utils.data import DataLoader +from utils import Logger, cal_loss, mkdir_p, progress_bar, save_args, save_model def parse_args(): """Parameters""" - parser = argparse.ArgumentParser('training') - parser.add_argument('-c', '--checkpoint', type=str, metavar='PATH', - help='path to save checkpoint (default: checkpoint)') - parser.add_argument('--msg', type=str, help='message after checkpoint') - parser.add_argument('--batch_size', type=int, default=32, help='batch size in training') - parser.add_argument('--model', default='PointNet', help='model name [default: pointnet_cls]') - parser.add_argument('--epoch', default=300, type=int, help='number of epoch in training') - parser.add_argument('--num_points', type=int, default=1024, help='Point Number') - parser.add_argument('--learning_rate', default=0.1, type=float, help='learning rate in training') - parser.add_argument('--min_lr', default=0.005, type=float, help='min lr') - parser.add_argument('--weight_decay', type=float, default=2e-4, help='decay rate') - parser.add_argument('--seed', type=int, help='random seed') - parser.add_argument('--workers', default=8, type=int, help='workers') + parser = argparse.ArgumentParser("training") + parser.add_argument( + "-c", + "--checkpoint", + type=str, + metavar="PATH", + help="path to save checkpoint (default: checkpoint)", + ) + parser.add_argument("--msg", type=str, help="message after checkpoint") + parser.add_argument("--batch_size", type=int, default=32, help="batch size in training") + parser.add_argument("--model", default="PointNet", help="model name [default: pointnet_cls]") + parser.add_argument("--epoch", default=300, type=int, help="number of epoch in training") + parser.add_argument("--num_points", type=int, default=1024, help="Point Number") + parser.add_argument("--learning_rate", default=0.1, type=float, help="learning rate in training") + parser.add_argument("--min_lr", default=0.005, type=float, help="min lr") + parser.add_argument("--weight_decay", type=float, default=2e-4, help="decay rate") + parser.add_argument("--seed", type=int, help="random seed") + parser.add_argument("--workers", default=8, type=int, help="workers") return parser.parse_args() @@ -46,7 +51,7 @@ def main(): os.environ["HDF5_USE_FILE_LOCKING"] = "FALSE" assert torch.cuda.is_available(), "Please ensure codes are executed in cuda." - device = 'cuda' + device = "cuda" if args.seed is not None: torch.manual_seed(args.seed) np.random.seed(args.seed) @@ -55,19 +60,19 @@ def main(): torch.set_printoptions(10) torch.backends.cudnn.benchmark = False torch.backends.cudnn.deterministic = True - os.environ['PYTHONHASHSEED'] = str(args.seed) - time_str = str(datetime.datetime.now().strftime('-%Y%m%d%H%M%S')) + os.environ["PYTHONHASHSEED"] = str(args.seed) + time_str = str(datetime.datetime.now().strftime("-%Y%m%d%H%M%S")) if args.msg is None: message = time_str else: message = "-" + args.msg - args.checkpoint = 'checkpoints/' + args.model + message + '-' + str(args.seed) + args.checkpoint = "checkpoints/" + args.model + message + "-" + str(args.seed) if not os.path.isdir(args.checkpoint): mkdir_p(args.checkpoint) screen_logger = logging.getLogger("Model") screen_logger.setLevel(logging.INFO) - formatter = logging.Formatter('%(message)s') + formatter = logging.Formatter("%(message)s") file_handler = logging.FileHandler(os.path.join(args.checkpoint, "out.txt")) file_handler.setLevel(logging.INFO) file_handler.setFormatter(formatter) @@ -79,19 +84,19 @@ def main(): # Model printf(f"args: {args}") - printf('==> Building model..') + printf("==> Building model..") net = models.__dict__[args.model]() criterion = cal_loss net = net.to(device) # criterion = criterion.to(device) - if device == 'cuda': + if device == "cuda": net = torch.nn.DataParallel(net) cudnn.benchmark = True - best_test_acc = 0. # best test accuracy - best_train_acc = 0. - best_test_acc_avg = 0. - best_train_acc_avg = 0. + best_test_acc = 0.0 # best test accuracy + best_train_acc = 0.0 + best_test_acc_avg = 0.0 + best_train_acc_avg = 0.0 best_test_loss = float("inf") best_train_loss = float("inf") start_epoch = 0 # start from epoch 0 or last checkpoint epoch @@ -99,30 +104,49 @@ def main(): if not os.path.isfile(os.path.join(args.checkpoint, "last_checkpoint.pth")): save_args(args) - logger = Logger(os.path.join(args.checkpoint, 'log.txt'), title="ModelNet" + args.model) - logger.set_names(["Epoch-Num", 'Learning-Rate', - 'Train-Loss', 'Train-acc-B', 'Train-acc', - 'Valid-Loss', 'Valid-acc-B', 'Valid-acc']) + logger = Logger(os.path.join(args.checkpoint, "log.txt"), title="ModelNet" + args.model) + logger.set_names( + [ + "Epoch-Num", + "Learning-Rate", + "Train-Loss", + "Train-acc-B", + "Train-acc", + "Valid-Loss", + "Valid-acc-B", + "Valid-acc", + ], + ) else: printf(f"Resuming last checkpoint from {args.checkpoint}") checkpoint_path = os.path.join(args.checkpoint, "last_checkpoint.pth") checkpoint = torch.load(checkpoint_path) - net.load_state_dict(checkpoint['net']) - start_epoch = checkpoint['epoch'] - best_test_acc = checkpoint['best_test_acc'] - best_train_acc = checkpoint['best_train_acc'] - best_test_acc_avg = checkpoint['best_test_acc_avg'] - best_train_acc_avg = checkpoint['best_train_acc_avg'] - best_test_loss = checkpoint['best_test_loss'] - best_train_loss = checkpoint['best_train_loss'] - logger = Logger(os.path.join(args.checkpoint, 'log.txt'), title="ModelNet" + args.model, resume=True) - optimizer_dict = checkpoint['optimizer'] + net.load_state_dict(checkpoint["net"]) + start_epoch = checkpoint["epoch"] + best_test_acc = checkpoint["best_test_acc"] + best_train_acc = checkpoint["best_train_acc"] + best_test_acc_avg = checkpoint["best_test_acc_avg"] + best_train_acc_avg = checkpoint["best_train_acc_avg"] + best_test_loss = checkpoint["best_test_loss"] + best_train_loss = checkpoint["best_train_loss"] + logger = Logger(os.path.join(args.checkpoint, "log.txt"), title="ModelNet" + args.model, resume=True) + optimizer_dict = checkpoint["optimizer"] - printf('==> Preparing data..') - train_loader = DataLoader(ModelNet40(partition='train', num_points=args.num_points), num_workers=args.workers, - batch_size=args.batch_size, shuffle=True, drop_last=True) - test_loader = DataLoader(ModelNet40(partition='test', num_points=args.num_points), num_workers=args.workers, - batch_size=args.batch_size // 2, shuffle=False, drop_last=False) + printf("==> Preparing data..") + train_loader = DataLoader( + ModelNet40(partition="train", num_points=args.num_points), + num_workers=args.workers, + batch_size=args.batch_size, + shuffle=True, + drop_last=True, + ) + test_loader = DataLoader( + ModelNet40(partition="test", num_points=args.num_points), + num_workers=args.workers, + batch_size=args.batch_size // 2, + shuffle=False, + drop_last=False, + ) optimizer = torch.optim.SGD(net.parameters(), lr=args.learning_rate, momentum=0.9, weight_decay=args.weight_decay) if optimizer_dict is not None: @@ -130,7 +154,7 @@ def main(): scheduler = CosineAnnealingLR(optimizer, args.epoch, eta_min=args.min_lr, last_epoch=start_epoch - 1) for epoch in range(start_epoch, args.epoch): - printf('Epoch(%d/%s) Learning Rate %s:' % (epoch + 1, args.epoch, optimizer.param_groups[0]['lr'])) + printf("Epoch(%d/%s) Learning Rate %s:" % (epoch + 1, args.epoch, optimizer.param_groups[0]["lr"])) train_out = train(net, train_loader, optimizer, criterion, device) # {"loss", "acc", "acc_avg", "time"} test_out = validate(net, test_loader, criterion, device) scheduler.step() @@ -149,31 +173,46 @@ def main(): best_train_loss = train_out["loss"] if (train_out["loss"] < best_train_loss) else best_train_loss save_model( - net, epoch, path=args.checkpoint, acc=test_out["acc"], is_best=is_best, + net, + epoch, + path=args.checkpoint, + acc=test_out["acc"], + is_best=is_best, best_test_acc=best_test_acc, # best test accuracy best_train_acc=best_train_acc, best_test_acc_avg=best_test_acc_avg, best_train_acc_avg=best_train_acc_avg, best_test_loss=best_test_loss, best_train_loss=best_train_loss, - optimizer=optimizer.state_dict() + optimizer=optimizer.state_dict(), + ) + logger.append( + [ + epoch, + optimizer.param_groups[0]["lr"], + train_out["loss"], + train_out["acc_avg"], + train_out["acc"], + test_out["loss"], + test_out["acc_avg"], + test_out["acc"], + ], ) - logger.append([epoch, optimizer.param_groups[0]['lr'], - train_out["loss"], train_out["acc_avg"], train_out["acc"], - test_out["loss"], test_out["acc_avg"], test_out["acc"]]) printf( - f"Training loss:{train_out['loss']} acc_avg:{train_out['acc_avg']}% acc:{train_out['acc']}% time:{train_out['time']}s") + f"Training loss:{train_out['loss']} acc_avg:{train_out['acc_avg']}% acc:{train_out['acc']}% time:{train_out['time']}s", + ) printf( f"Testing loss:{test_out['loss']} acc_avg:{test_out['acc_avg']}% " - f"acc:{test_out['acc']}% time:{test_out['time']}s [best test acc: {best_test_acc}%] \n\n") + f"acc:{test_out['acc']}% time:{test_out['time']}s [best test acc: {best_test_acc}%] \n\n", + ) logger.close() - printf(f"++++++++" * 2 + "Final results" + "++++++++" * 2) + printf("++++++++" * 2 + "Final results" + "++++++++" * 2) printf(f"++ Last Train time: {train_out['time']} | Last Test time: {test_out['time']} ++") printf(f"++ Best Train loss: {best_train_loss} | Best Test loss: {best_test_loss} ++") printf(f"++ Best Train acc_B: {best_train_acc_avg} | Best Test acc_B: {best_test_acc_avg} ++") printf(f"++ Best Train acc: {best_train_acc} | Best Test acc: {best_test_acc} ++") - printf(f"++++++++" * 5) + printf("++++++++" * 5) def train(net, trainloader, optimizer, criterion, device): @@ -202,17 +241,21 @@ def train(net, trainloader, optimizer, criterion, device): total += label.size(0) correct += preds.eq(label).sum().item() - progress_bar(batch_idx, len(trainloader), 'Loss: %.3f | Acc: %.3f%% (%d/%d)' - % (train_loss / (batch_idx + 1), 100. * correct / total, correct, total)) + progress_bar( + batch_idx, + len(trainloader), + "Loss: %.3f | Acc: %.3f%% (%d/%d)" + % (train_loss / (batch_idx + 1), 100.0 * correct / total, correct, total), + ) time_cost = int((datetime.datetime.now() - time_cost).total_seconds()) train_true = np.concatenate(train_true) train_pred = np.concatenate(train_pred) return { "loss": float("%.3f" % (train_loss / (batch_idx + 1))), - "acc": float("%.3f" % (100. * metrics.accuracy_score(train_true, train_pred))), - "acc_avg": float("%.3f" % (100. * metrics.balanced_accuracy_score(train_true, train_pred))), - "time": time_cost + "acc": float("%.3f" % (100.0 * metrics.accuracy_score(train_true, train_pred))), + "acc_avg": float("%.3f" % (100.0 * metrics.balanced_accuracy_score(train_true, train_pred))), + "time": time_cost, } @@ -236,19 +279,23 @@ def validate(net, testloader, criterion, device): test_pred.append(preds.detach().cpu().numpy()) total += label.size(0) correct += preds.eq(label).sum().item() - progress_bar(batch_idx, len(testloader), 'Loss: %.3f | Acc: %.3f%% (%d/%d)' - % (test_loss / (batch_idx + 1), 100. * correct / total, correct, total)) + progress_bar( + batch_idx, + len(testloader), + "Loss: %.3f | Acc: %.3f%% (%d/%d)" + % (test_loss / (batch_idx + 1), 100.0 * correct / total, correct, total), + ) time_cost = int((datetime.datetime.now() - time_cost).total_seconds()) test_true = np.concatenate(test_true) test_pred = np.concatenate(test_pred) return { "loss": float("%.3f" % (test_loss / (batch_idx + 1))), - "acc": float("%.3f" % (100. * metrics.accuracy_score(test_true, test_pred))), - "acc_avg": float("%.3f" % (100. * metrics.balanced_accuracy_score(test_true, test_pred))), - "time": time_cost + "acc": float("%.3f" % (100.0 * metrics.accuracy_score(test_true, test_pred))), + "acc_avg": float("%.3f" % (100.0 * metrics.balanced_accuracy_score(test_true, test_pred))), + "time": time_cost, } -if __name__ == '__main__': +if __name__ == "__main__": main() diff --git a/classification_ModelNet40/models/__init__.py b/classification_ModelNet40/models/__init__.py index b8815da..5a546cf 100644 --- a/classification_ModelNet40/models/__init__.py +++ b/classification_ModelNet40/models/__init__.py @@ -1,3 +1 @@ -from __future__ import absolute_import - from .pointmlp import pointMLP, pointMLPElite diff --git a/classification_ModelNet40/models/pointmlp.py b/classification_ModelNet40/models/pointmlp.py index d350efb..236db49 100644 --- a/classification_ModelNet40/models/pointmlp.py +++ b/classification_ModelNet40/models/pointmlp.py @@ -1,35 +1,32 @@ - import torch import torch.nn as nn import torch.nn.functional as F + # from torch import einsum # from einops import rearrange, repeat - - from pointnet2_ops import pointnet2_utils def get_activation(activation): - if activation.lower() == 'gelu': + if activation.lower() == "gelu": return nn.GELU() - elif activation.lower() == 'rrelu': + elif activation.lower() == "rrelu": return nn.RReLU(inplace=True) - elif activation.lower() == 'selu': + elif activation.lower() == "selu": return nn.SELU(inplace=True) - elif activation.lower() == 'silu': + elif activation.lower() == "silu": return nn.SiLU(inplace=True) - elif activation.lower() == 'hardswish': + elif activation.lower() == "hardswish": return nn.Hardswish(inplace=True) - elif activation.lower() == 'leakyrelu': + elif activation.lower() == "leakyrelu": return nn.LeakyReLU(inplace=True) else: return nn.ReLU(inplace=True) def square_distance(src, dst): - """ - Calculate Euclid distance between each two points. - src^T * dst = xn * xm + yn * ym + zn * zm； + """Calculate Euclid distance between each two points. + src^T * dst = xn * xm + yn * ym + zn * zm; sum(src^2, dim=-1) = xn*xn + yn*yn + zn*zn; sum(dst^2, dim=-1) = xm*xm + ym*ym + zm*zm; dist = (xn - xm)^2 + (yn - ym)^2 + (zn - zm)^2 @@ -38,23 +35,23 @@ def square_distance(src, dst): src: source points, [B, N, C] dst: target points, [B, M, C] Output: - dist: per-point square distance, [B, N, M] + dist: per-point square distance, [B, N, M]. """ B, N, _ = src.shape _, M, _ = dst.shape dist = -2 * torch.matmul(src, dst.permute(0, 2, 1)) - dist += torch.sum(src ** 2, -1).view(B, N, 1) - dist += torch.sum(dst ** 2, -1).view(B, 1, M) + dist += torch.sum(src**2, -1).view(B, N, 1) + dist += torch.sum(dst**2, -1).view(B, 1, M) return dist def index_points(points, idx): - """ - Input: + """Input: points: input points data, [B, N, C] - idx: sample index data, [B, S] + idx: sample index data, [B, S]. + Return: - new_points:, indexed points data, [B, S, C] + new_points:, indexed points data, [B, S, C]. """ device = points.device B = points.shape[0] @@ -63,17 +60,15 @@ def index_points(points, idx): repeat_shape = list(idx.shape) repeat_shape[0] = 1 batch_indices = torch.arange(B, dtype=torch.long).to(device).view(view_shape).repeat(repeat_shape) - new_points = points[batch_indices, idx, :] - return new_points + return points[batch_indices, idx, :] def farthest_point_sample(xyz, npoint): - """ - Input: + """Input: xyz: pointcloud data, [B, N, 3] npoint: number of samples Return: - centroids: sampled pointcloud index, [B, npoint] + centroids: sampled pointcloud index, [B, npoint]. """ device = xyz.device B, N, C = xyz.shape @@ -91,21 +86,21 @@ def farthest_point_sample(xyz, npoint): def query_ball_point(radius, nsample, xyz, new_xyz): - """ - Input: + """Input: radius: local region radius nsample: max sample number in local region xyz: all points, [B, N, 3] - new_xyz: query points, [B, S, 3] + new_xyz: query points, [B, S, 3]. + Return: - group_idx: grouped points index, [B, S, nsample] + group_idx: grouped points index, [B, S, nsample]. """ device = xyz.device B, N, C = xyz.shape _, S, _ = new_xyz.shape group_idx = torch.arange(N, dtype=torch.long).to(device).view(1, 1, N).repeat([B, S, 1]) sqrdists = square_distance(new_xyz, xyz) - group_idx[sqrdists > radius ** 2] = N + group_idx[sqrdists > radius**2] = N group_idx = group_idx.sort(dim=-1)[0][:, :, :nsample] group_first = group_idx[:, :, 0].view(B, S, 1).repeat([1, 1, nsample]) mask = group_idx == N @@ -114,13 +109,13 @@ def query_ball_point(radius, nsample, xyz, new_xyz): def knn_point(nsample, xyz, new_xyz): - """ - Input: + """Input: nsample: max sample number in local region xyz: all points, [B, N, C] - new_xyz: query points, [B, S, C] + new_xyz: query points, [B, S, C]. + Return: - group_idx: grouped points index, [B, S, nsample] + group_idx: grouped points index, [B, S, nsample]. """ sqrdists = square_distance(new_xyz, xyz) _, group_idx = torch.topk(sqrdists, nsample, dim=-1, largest=False, sorted=False) @@ -129,13 +124,12 @@ def knn_point(nsample, xyz, new_xyz): class LocalGrouper(nn.Module): def __init__(self, channel, groups, kneighbors, use_xyz=True, normalize="center", **kwargs): - """ - Give xyz[b,p,3] and fea[b,p,d], return new_xyz[b,g,3] and new_fea[b,g,k,d] + """Give xyz[b,p,3] and fea[b,p,d], return new_xyz[b,g,3] and new_fea[b,g,k,d] :param groups: groups number :param kneighbors: k-nerighbors - :param kwargs: others + :param kwargs: others. """ - super(LocalGrouper, self).__init__() + super().__init__() self.groups = groups self.kneighbors = kneighbors self.use_xyz = use_xyz @@ -144,11 +138,11 @@ class LocalGrouper(nn.Module): else: self.normalize = None if self.normalize not in ["center", "anchor"]: - print(f"Unrecognized normalize parameter (self.normalize), set to None. Should be one of [center, anchor].") + print("Unrecognized normalize parameter (self.normalize), set to None. Should be one of [center, anchor].") self.normalize = None if self.normalize is not None: - add_channel=3 if self.use_xyz else 0 - self.affine_alpha = nn.Parameter(torch.ones([1,1,1,channel + add_channel])) + add_channel = 3 if self.use_xyz else 0 + self.affine_alpha = nn.Parameter(torch.ones([1, 1, 1, channel + add_channel])) self.affine_beta = nn.Parameter(torch.zeros([1, 1, 1, channel + add_channel])) def forward(self, xyz, points): @@ -167,29 +161,33 @@ class LocalGrouper(nn.Module): grouped_xyz = index_points(xyz, idx) # [B, npoint, k, 3] grouped_points = index_points(points, idx) # [B, npoint, k, d] if self.use_xyz: - grouped_points = torch.cat([grouped_points, grouped_xyz],dim=-1) # [B, npoint, k, d+3] + grouped_points = torch.cat([grouped_points, grouped_xyz], dim=-1) # [B, npoint, k, d+3] if self.normalize is not None: - if self.normalize =="center": + if self.normalize == "center": mean = torch.mean(grouped_points, dim=2, keepdim=True) - if self.normalize =="anchor": - mean = torch.cat([new_points, new_xyz],dim=-1) if self.use_xyz else new_points + if self.normalize == "anchor": + mean = torch.cat([new_points, new_xyz], dim=-1) if self.use_xyz else new_points mean = mean.unsqueeze(dim=-2) # [B, npoint, 1, d+3] - std = torch.std((grouped_points-mean).reshape(B,-1),dim=-1,keepdim=True).unsqueeze(dim=-1).unsqueeze(dim=-1) - grouped_points = (grouped_points-mean)/(std + 1e-5) - grouped_points = self.affine_alpha*grouped_points + self.affine_beta + std = ( + torch.std((grouped_points - mean).reshape(B, -1), dim=-1, keepdim=True) + .unsqueeze(dim=-1) + .unsqueeze(dim=-1) + ) + grouped_points = (grouped_points - mean) / (std + 1e-5) + grouped_points = self.affine_alpha * grouped_points + self.affine_beta new_points = torch.cat([grouped_points, new_points.view(B, S, 1, -1).repeat(1, 1, self.kneighbors, 1)], dim=-1) return new_xyz, new_points class ConvBNReLU1D(nn.Module): - def __init__(self, in_channels, out_channels, kernel_size=1, bias=True, activation='relu'): - super(ConvBNReLU1D, self).__init__() + def __init__(self, in_channels, out_channels, kernel_size=1, bias=True, activation="relu"): + super().__init__() self.act = get_activation(activation) self.net = nn.Sequential( nn.Conv1d(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, bias=bias), nn.BatchNorm1d(out_channels), - self.act + self.act, ) def forward(self, x): @@ -197,30 +195,43 @@ class ConvBNReLU1D(nn.Module): class ConvBNReLURes1D(nn.Module): - def __init__(self, channel, kernel_size=1, groups=1, res_expansion=1.0, bias=True, activation='relu'): - super(ConvBNReLURes1D, self).__init__() + def __init__(self, channel, kernel_size=1, groups=1, res_expansion=1.0, bias=True, activation="relu"): + super().__init__() self.act = get_activation(activation) self.net1 = nn.Sequential( - nn.Conv1d(in_channels=channel, out_channels=int(channel * res_expansion), - kernel_size=kernel_size, groups=groups, bias=bias), + nn.Conv1d( + in_channels=channel, + out_channels=int(channel * res_expansion), + kernel_size=kernel_size, + groups=groups, + bias=bias, + ), nn.BatchNorm1d(int(channel * res_expansion)), - self.act + self.act, ) if groups > 1: self.net2 = nn.Sequential( - nn.Conv1d(in_channels=int(channel * res_expansion), out_channels=channel, - kernel_size=kernel_size, groups=groups, bias=bias), + nn.Conv1d( + in_channels=int(channel * res_expansion), + out_channels=channel, + kernel_size=kernel_size, + groups=groups, + bias=bias, + ), nn.BatchNorm1d(channel), self.act, - nn.Conv1d(in_channels=channel, out_channels=channel, - kernel_size=kernel_size, bias=bias), + nn.Conv1d(in_channels=channel, out_channels=channel, kernel_size=kernel_size, bias=bias), nn.BatchNorm1d(channel), ) else: self.net2 = nn.Sequential( - nn.Conv1d(in_channels=int(channel * res_expansion), out_channels=channel, - kernel_size=kernel_size, bias=bias), - nn.BatchNorm1d(channel) + nn.Conv1d( + in_channels=int(channel * res_expansion), + out_channels=channel, + kernel_size=kernel_size, + bias=bias, + ), + nn.BatchNorm1d(channel), ) def forward(self, x): @@ -228,21 +239,34 @@ class ConvBNReLURes1D(nn.Module): class PreExtraction(nn.Module): - def __init__(self, channels, out_channels, blocks=1, groups=1, res_expansion=1, bias=True, - activation='relu', use_xyz=True): - """ - input: [b,g,k,d]: output:[b,d,g] + def __init__( + self, + channels, + out_channels, + blocks=1, + groups=1, + res_expansion=1, + bias=True, + activation="relu", + use_xyz=True, + ): + """input: [b,g,k,d]: output:[b,d,g] :param channels: :param blocks: """ - super(PreExtraction, self).__init__() - in_channels = 3+2*channels if use_xyz else 2*channels + super().__init__() + in_channels = 3 + 2 * channels if use_xyz else 2 * channels self.transfer = ConvBNReLU1D(in_channels, out_channels, bias=bias, activation=activation) operation = [] for _ in range(blocks): operation.append( - ConvBNReLURes1D(out_channels, groups=groups, res_expansion=res_expansion, - bias=bias, activation=activation) + ConvBNReLURes1D( + out_channels, + groups=groups, + res_expansion=res_expansion, + bias=bias, + activation=activation, + ), ) self.operation = nn.Sequential(*operation) @@ -254,22 +278,20 @@ class PreExtraction(nn.Module): batch_size, _, _ = x.size() x = self.operation(x) # [b, d, k] x = F.adaptive_max_pool1d(x, 1).view(batch_size, -1) - x = x.reshape(b, n, -1).permute(0, 2, 1) - return x + return x.reshape(b, n, -1).permute(0, 2, 1) class PosExtraction(nn.Module): - def __init__(self, channels, blocks=1, groups=1, res_expansion=1, bias=True, activation='relu'): - """ - input[b,d,g]; output[b,d,g] + def __init__(self, channels, blocks=1, groups=1, res_expansion=1, bias=True, activation="relu"): + """input[b,d,g]; output[b,d,g] :param channels: :param blocks: """ - super(PosExtraction, self).__init__() + super().__init__() operation = [] for _ in range(blocks): operation.append( - ConvBNReLURes1D(channels, groups=groups, res_expansion=res_expansion, bias=bias, activation=activation) + ConvBNReLURes1D(channels, groups=groups, res_expansion=res_expansion, bias=bias, activation=activation), ) self.operation = nn.Sequential(*operation) @@ -278,17 +300,32 @@ class PosExtraction(nn.Module): class Model(nn.Module): - def __init__(self, points=1024, class_num=40, embed_dim=64, groups=1, res_expansion=1.0, - activation="relu", bias=True, use_xyz=True, normalize="center", - dim_expansion=[2, 2, 2, 2], pre_blocks=[2, 2, 2, 2], pos_blocks=[2, 2, 2, 2], - k_neighbors=[32, 32, 32, 32], reducers=[2, 2, 2, 2], **kwargs): - super(Model, self).__init__() + def __init__( + self, + points=1024, + class_num=40, + embed_dim=64, + groups=1, + res_expansion=1.0, + activation="relu", + bias=True, + use_xyz=True, + normalize="center", + dim_expansion=[2, 2, 2, 2], + pre_blocks=[2, 2, 2, 2], + pos_blocks=[2, 2, 2, 2], + k_neighbors=[32, 32, 32, 32], + reducers=[2, 2, 2, 2], + **kwargs, + ): + super().__init__() self.stages = len(pre_blocks) self.class_num = class_num self.points = points self.embedding = ConvBNReLU1D(3, embed_dim, bias=bias, activation=activation) - assert len(pre_blocks) == len(k_neighbors) == len(reducers) == len(pos_blocks) == len(dim_expansion), \ - "Please check stage number consistent for pre_blocks, pos_blocks k_neighbors, reducers." + assert ( + len(pre_blocks) == len(k_neighbors) == len(reducers) == len(pos_blocks) == len(dim_expansion) + ), "Please check stage number consistent for pre_blocks, pos_blocks k_neighbors, reducers." self.local_grouper_list = nn.ModuleList() self.pre_blocks_list = nn.ModuleList() self.pos_blocks_list = nn.ModuleList() @@ -305,13 +342,26 @@ class Model(nn.Module): local_grouper = LocalGrouper(last_channel, anchor_points, kneighbor, use_xyz, normalize) # [b,g,k,d] self.local_grouper_list.append(local_grouper) # append pre_block_list - pre_block_module = PreExtraction(last_channel, out_channel, pre_block_num, groups=groups, - res_expansion=res_expansion, - bias=bias, activation=activation, use_xyz=use_xyz) + pre_block_module = PreExtraction( + last_channel, + out_channel, + pre_block_num, + groups=groups, + res_expansion=res_expansion, + bias=bias, + activation=activation, + use_xyz=use_xyz, + ) self.pre_blocks_list.append(pre_block_module) # append pos_block_list - pos_block_module = PosExtraction(out_channel, pos_block_num, groups=groups, - res_expansion=res_expansion, bias=bias, activation=activation) + pos_block_module = PosExtraction( + out_channel, + pos_block_num, + groups=groups, + res_expansion=res_expansion, + bias=bias, + activation=activation, + ) self.pos_blocks_list.append(pos_block_module) last_channel = out_channel @@ -326,7 +376,7 @@ class Model(nn.Module): nn.BatchNorm1d(256), self.act, nn.Dropout(0.5), - nn.Linear(256, self.class_num) + nn.Linear(256, self.class_num), ) def forward(self, x): @@ -340,26 +390,50 @@ class Model(nn.Module): x = self.pos_blocks_list[i](x) # [b,d,g] x = F.adaptive_max_pool1d(x, 1).squeeze(dim=-1) - x = self.classifier(x) - return x - - + return self.classifier(x) def pointMLP(num_classes=40, **kwargs) -> Model: - return Model(points=1024, class_num=num_classes, embed_dim=64, groups=1, res_expansion=1.0, - activation="relu", bias=False, use_xyz=False, normalize="anchor", - dim_expansion=[2, 2, 2, 2], pre_blocks=[2, 2, 2, 2], pos_blocks=[2, 2, 2, 2], - k_neighbors=[24, 24, 24, 24], reducers=[2, 2, 2, 2], **kwargs) + return Model( + points=1024, + class_num=num_classes, + embed_dim=64, + groups=1, + res_expansion=1.0, + activation="relu", + bias=False, + use_xyz=False, + normalize="anchor", + dim_expansion=[2, 2, 2, 2], + pre_blocks=[2, 2, 2, 2], + pos_blocks=[2, 2, 2, 2], + k_neighbors=[24, 24, 24, 24], + reducers=[2, 2, 2, 2], + **kwargs, + ) def pointMLPElite(num_classes=40, **kwargs) -> Model: - return Model(points=1024, class_num=num_classes, embed_dim=32, groups=1, res_expansion=0.25, - activation="relu", bias=False, use_xyz=False, normalize="anchor", - dim_expansion=[2, 2, 2, 1], pre_blocks=[1, 1, 2, 1], pos_blocks=[1, 1, 2, 1], - k_neighbors=[24,24,24,24], reducers=[2, 2, 2, 2], **kwargs) + return Model( + points=1024, + class_num=num_classes, + embed_dim=32, + groups=1, + res_expansion=0.25, + activation="relu", + bias=False, + use_xyz=False, + normalize="anchor", + dim_expansion=[2, 2, 2, 1], + pre_blocks=[1, 1, 2, 1], + pos_blocks=[1, 1, 2, 1], + k_neighbors=[24, 24, 24, 24], + reducers=[2, 2, 2, 2], + **kwargs, + ) -if __name__ == '__main__': + +if __name__ == "__main__": data = torch.rand(2, 3, 1024).cuda() print(data.shape) diff --git a/classification_ModelNet40/test.py b/classification_ModelNet40/test.py index 698cc46..35a1e37 100644 --- a/classification_ModelNet40/test.py +++ b/classification_ModelNet40/test.py @@ -1,49 +1,52 @@ -""" -python test.py --model pointMLP --msg 20220209053148-404 -""" +"""python test.py --model pointMLP --msg 20220209053148-404.""" import argparse -import os import datetime +import os + +import models as models +import numpy as np +import sklearn.metrics as metrics import torch -import torch.nn.parallel import torch.backends.cudnn as cudnn +import torch.nn.parallel import torch.optim import torch.utils.data import torch.utils.data.distributed -from torch.utils.data import DataLoader -import models as models -from utils import progress_bar, IOStream from data import ModelNet40 -import sklearn.metrics as metrics from helper import cal_loss -import numpy as np -import torch.nn.functional as F +from torch.utils.data import DataLoader +from utils import progress_bar -model_names = sorted(name for name in models.__dict__ - if callable(models.__dict__[name])) +model_names = sorted(name for name in models.__dict__ if callable(models.__dict__[name])) def parse_args(): """Parameters""" - parser = argparse.ArgumentParser('training') - parser.add_argument('-c', '--checkpoint', type=str, metavar='PATH', - help='path to save checkpoint (default: checkpoint)') - parser.add_argument('--msg', type=str, help='message after checkpoint') - parser.add_argument('--batch_size', type=int, default=16, help='batch size in training') - parser.add_argument('--model', default='pointMLP', help='model name [default: pointnet_cls]') - parser.add_argument('--num_classes', default=40, type=int, choices=[10, 40], help='training on ModelNet10/40') - parser.add_argument('--num_points', type=int, default=1024, help='Point Number') + parser = argparse.ArgumentParser("training") + parser.add_argument( + "-c", + "--checkpoint", + type=str, + metavar="PATH", + help="path to save checkpoint (default: checkpoint)", + ) + parser.add_argument("--msg", type=str, help="message after checkpoint") + parser.add_argument("--batch_size", type=int, default=16, help="batch size in training") + parser.add_argument("--model", default="pointMLP", help="model name [default: pointnet_cls]") + parser.add_argument("--num_classes", default=40, type=int, choices=[10, 40], help="training on ModelNet10/40") + parser.add_argument("--num_points", type=int, default=1024, help="Point Number") return parser.parse_args() + def main(): args = parse_args() print(f"args: {args}") os.environ["HDF5_USE_FILE_LOCKING"] = "FALSE" if torch.cuda.is_available(): - device = 'cuda' + device = "cuda" else: - device = 'cpu' + device = "cpu" print(f"==> Using device: {device}") # if args.msg is None: # message = str(datetime.datetime.now().strftime('-%Y%m%d%H%M%S')) @@ -53,26 +56,26 @@ def main(): if args.checkpoint is not None: print(f"==> Using checkpoint: {args.checkpoint}") - print('==> Preparing data..') + print("==> Preparing data..") test_loader = DataLoader( - ModelNet40(partition='test', num_points=args.num_points), + ModelNet40(partition="test", num_points=args.num_points), num_workers=4, batch_size=args.batch_size, shuffle=False, - drop_last=False + drop_last=False, ) # Model - print('==> Building model..') + print("==> Building model..") net = models.__dict__[args.model]() criterion = cal_loss net = net.to(device) # checkpoint_path = os.path.join(args.checkpoint, 'best_checkpoint.pth') - checkpoint = torch.load(args.checkpoint, map_location=torch.device('cpu')) + checkpoint = torch.load(args.checkpoint, map_location=torch.device("cpu")) # criterion = criterion.to(device) - if device == 'cuda': + if device == "cuda": net = torch.nn.DataParallel(net) cudnn.benchmark = True - net.load_state_dict(checkpoint['net']) + net.load_state_dict(checkpoint["net"]) test_out = validate(net, test_loader, criterion, device) print(f"Vanilla out: {test_out}") @@ -98,19 +101,23 @@ def validate(net, testloader, criterion, device): test_pred.append(preds.detach().cpu().numpy()) total += label.size(0) correct += preds.eq(label).sum().item() - progress_bar(batch_idx, len(testloader), 'Loss: %.3f | Acc: %.3f%% (%d/%d)' - % (test_loss / (batch_idx + 1), 100. * correct / total, correct, total)) + progress_bar( + batch_idx, + len(testloader), + "Loss: %.3f | Acc: %.3f%% (%d/%d)" + % (test_loss / (batch_idx + 1), 100.0 * correct / total, correct, total), + ) time_cost = int((datetime.datetime.now() - time_cost).total_seconds()) test_true = np.concatenate(test_true) test_pred = np.concatenate(test_pred) return { "loss": float("%.3f" % (test_loss / (batch_idx + 1))), - "acc": float("%.3f" % (100. * metrics.accuracy_score(test_true, test_pred))), - "acc_avg": float("%.3f" % (100. * metrics.balanced_accuracy_score(test_true, test_pred))), - "time": time_cost + "acc": float("%.3f" % (100.0 * metrics.accuracy_score(test_true, test_pred))), + "acc_avg": float("%.3f" % (100.0 * metrics.balanced_accuracy_score(test_true, test_pred))), + "time": time_cost, } -if __name__ == '__main__': +if __name__ == "__main__": main() diff --git a/classification_ModelNet40/utils/__init__.py b/classification_ModelNet40/utils/__init__.py index 09e4b76..9be60ee 100644 --- a/classification_ModelNet40/utils/__init__.py +++ b/classification_ModelNet40/utils/__init__.py @@ -1,5 +1,4 @@ -"""Useful utils -""" -from .misc import * +"""Useful utils.""" from .logger import * +from .misc import * from .progress.progress.bar import Bar as Bar diff --git a/classification_ModelNet40/utils/logger.py b/classification_ModelNet40/utils/logger.py index 7eb5c67..db15f5f 100644 --- a/classification_ModelNet40/utils/logger.py +++ b/classification_ModelNet40/utils/logger.py @@ -1,89 +1,92 @@ # A simple torch style logger # (C) Wei YANG 2017 -from __future__ import absolute_import + import matplotlib.pyplot as plt -import os -import sys import numpy as np -__all__ = ['Logger', 'LoggerMonitor', 'savefig'] +__all__ = ["Logger", "LoggerMonitor", "savefig"] + def savefig(fname, dpi=None): - dpi = 150 if dpi == None else dpi + dpi = 150 if dpi is None else dpi plt.savefig(fname, dpi=dpi) - + + def plot_overlap(logger, names=None): - names = logger.names if names == None else names + names = logger.names if names is None else names numbers = logger.numbers for _, name in enumerate(names): x = np.arange(len(numbers[name])) plt.plot(x, np.asarray(numbers[name])) - return [logger.title + '(' + name + ')' for name in names] + return [logger.title + "(" + name + ")" for name in names] -class Logger(object): - '''Save training process to log file with simple plot function.''' - def __init__(self, fpath, title=None, resume=False): + +class Logger: + """Save training process to log file with simple plot function.""" + + def __init__(self, fpath, title=None, resume=False): self.file = None self.resume = resume - self.title = '' if title == None else title + self.title = "" if title is None else title if fpath is not None: - if resume: - self.file = open(fpath, 'r') + if resume: + self.file = open(fpath) name = self.file.readline() - self.names = name.rstrip().split('\t') + self.names = name.rstrip().split("\t") self.numbers = {} for _, name in enumerate(self.names): self.numbers[name] = [] for numbers in self.file: - numbers = numbers.rstrip().split('\t') + numbers = numbers.rstrip().split("\t") for i in range(0, len(numbers)): self.numbers[self.names[i]].append(numbers[i]) self.file.close() - self.file = open(fpath, 'a') + self.file = open(fpath, "a") else: - self.file = open(fpath, 'w') + self.file = open(fpath, "w") def set_names(self, names): - if self.resume: + if self.resume: pass # initialize numbers as empty list self.numbers = {} self.names = names for _, name in enumerate(self.names): self.file.write(name) - self.file.write('\t') + self.file.write("\t") self.numbers[name] = [] - self.file.write('\n') + self.file.write("\n") self.file.flush() - def append(self, numbers): - assert len(self.names) == len(numbers), 'Numbers do not match names' + assert len(self.names) == len(numbers), "Numbers do not match names" for index, num in enumerate(numbers): - self.file.write("{0:.6f}".format(num)) - self.file.write('\t') + self.file.write(f"{num:.6f}") + self.file.write("\t") self.numbers[self.names[index]].append(num) - self.file.write('\n') + self.file.write("\n") self.file.flush() - def plot(self, names=None): - names = self.names if names == None else names + def plot(self, names=None): + names = self.names if names is None else names numbers = self.numbers for _, name in enumerate(names): x = np.arange(len(numbers[name])) plt.plot(x, np.asarray(numbers[name])) - plt.legend([self.title + '(' + name + ')' for name in names]) + plt.legend([self.title + "(" + name + ")" for name in names]) plt.grid(True) def close(self): if self.file is not None: self.file.close() -class LoggerMonitor(object): - '''Load and visualize multiple logs.''' - def __init__ (self, paths): - '''paths is a distionary with {name:filepath} pair''' + +class LoggerMonitor: + """Load and visualize multiple logs.""" + + def __init__(self, paths): + """Paths is a distionary with {name:filepath} pair.""" self.loggers = [] for title, path in paths.items(): logger = Logger(path, title=title, resume=True) @@ -95,10 +98,11 @@ class LoggerMonitor(object): legend_text = [] for logger in self.loggers: legend_text += plot_overlap(logger, names) - plt.legend(legend_text, bbox_to_anchor=(1.05, 1), loc=2, borderaxespad=0.) + plt.legend(legend_text, bbox_to_anchor=(1.05, 1), loc=2, borderaxespad=0.0) plt.grid(True) - -if __name__ == '__main__': + + +if __name__ == "__main__": # # Example # logger = Logger('test.txt') # logger.set_names(['Train loss', 'Valid loss','Test loss']) @@ -115,13 +119,13 @@ if __name__ == '__main__': # Example: logger monitor paths = { - 'resadvnet20':'/home/wyang/code/pytorch-classification/checkpoint/cifar10/resadvnet20/log.txt', - 'resadvnet32':'/home/wyang/code/pytorch-classification/checkpoint/cifar10/resadvnet32/log.txt', - 'resadvnet44':'/home/wyang/code/pytorch-classification/checkpoint/cifar10/resadvnet44/log.txt', + "resadvnet20": "/home/wyang/code/pytorch-classification/checkpoint/cifar10/resadvnet20/log.txt", + "resadvnet32": "/home/wyang/code/pytorch-classification/checkpoint/cifar10/resadvnet32/log.txt", + "resadvnet44": "/home/wyang/code/pytorch-classification/checkpoint/cifar10/resadvnet44/log.txt", } - field = ['Valid Acc.'] + field = ["Valid Acc."] monitor = LoggerMonitor(paths) monitor.plot(names=field) - savefig('test.eps') \ No newline at end of file + savefig("test.eps") diff --git a/classification_ModelNet40/utils/misc.py b/classification_ModelNet40/utils/misc.py index 35f743a..3d073b8 100644 --- a/classification_ModelNet40/utils/misc.py +++ b/classification_ModelNet40/utils/misc.py @@ -1,48 +1,56 @@ -'''Some helper functions for PyTorch, including: - - get_mean_and_std: calculate the mean and std value of dataset. - - msr_init: net parameter initialization. - - progress_bar: progress bar mimic xlua.progress. -''' +"""Some helper functions for PyTorch, including: +- get_mean_and_std: calculate the mean and std value of dataset. +- msr_init: net parameter initialization. +- progress_bar: progress bar mimic xlua.progress. +""" import errno import os +import random +import shutil import sys import time -import math -import torch -import shutil + import numpy as np -import random -import torch.nn.functional as F - - +import torch import torch.nn as nn +import torch.nn.functional as F import torch.nn.init as init -from torch.autograd import Variable -__all__ = ['get_mean_and_std', 'init_params', 'mkdir_p', 'AverageMeter', - 'progress_bar','save_model',"save_args","set_seed", "IOStream", "cal_loss"] +__all__ = [ + "get_mean_and_std", + "init_params", + "mkdir_p", + "AverageMeter", + "progress_bar", + "save_model", + "save_args", + "set_seed", + "IOStream", + "cal_loss", +] def get_mean_and_std(dataset): - '''Compute the mean and std value of dataset.''' - dataloader = trainloader = torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=True, num_workers=2) + """Compute the mean and std value of dataset.""" + dataloader = torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=True, num_workers=2) mean = torch.zeros(3) std = torch.zeros(3) - print('==> Computing mean and std..') - for inputs, targets in dataloader: + print("==> Computing mean and std..") + for inputs, _targets in dataloader: for i in range(3): - mean[i] += inputs[:,i,:,:].mean() - std[i] += inputs[:,i,:,:].std() + mean[i] += inputs[:, i, :, :].mean() + std[i] += inputs[:, i, :, :].std() mean.div_(len(dataset)) std.div_(len(dataset)) return mean, std + def init_params(net): - '''Init layer parameters.''' + """Init layer parameters.""" for m in net.modules(): if isinstance(m, nn.Conv2d): - init.kaiming_normal(m.weight, mode='fan_out') + init.kaiming_normal(m.weight, mode="fan_out") if m.bias: init.constant(m.bias, 0) elif isinstance(m, nn.BatchNorm2d): @@ -53,8 +61,9 @@ def init_params(net): if m.bias: init.constant(m.bias, 0) + def mkdir_p(path): - '''make dir if not exist''' + """Make dir if not exist.""" try: os.makedirs(path) except OSError as exc: # Python >2.5 @@ -63,10 +72,12 @@ def mkdir_p(path): else: raise -class AverageMeter(object): + +class AverageMeter: """Computes and stores the average and current value - Imported from https://github.com/pytorch/examples/blob/master/imagenet/main.py#L247-L262 + Imported from https://github.com/pytorch/examples/blob/master/imagenet/main.py#L247-L262. """ + def __init__(self): self.reset() @@ -83,25 +94,26 @@ class AverageMeter(object): self.avg = self.sum / self.count - -TOTAL_BAR_LENGTH = 65. +TOTAL_BAR_LENGTH = 65.0 last_time = time.time() begin_time = last_time + + def progress_bar(current, total, msg=None): global last_time, begin_time if current == 0: begin_time = time.time() # Reset for new bar. - cur_len = int(TOTAL_BAR_LENGTH*current/total) + cur_len = int(TOTAL_BAR_LENGTH * current / total) rest_len = int(TOTAL_BAR_LENGTH - cur_len) - 1 - sys.stdout.write(' [') - for i in range(cur_len): - sys.stdout.write('=') - sys.stdout.write('>') - for i in range(rest_len): - sys.stdout.write('.') - sys.stdout.write(']') + sys.stdout.write(" [") + for _i in range(cur_len): + sys.stdout.write("=") + sys.stdout.write(">") + for _i in range(rest_len): + sys.stdout.write(".") + sys.stdout.write("]") cur_time = time.time() step_time = cur_time - last_time @@ -109,12 +121,12 @@ def progress_bar(current, total, msg=None): tot_time = cur_time - begin_time L = [] - L.append(' Step: %s' % format_time(step_time)) - L.append(' | Tot: %s' % format_time(tot_time)) + L.append(" Step: %s" % format_time(step_time)) + L.append(" | Tot: %s" % format_time(tot_time)) if msg: - L.append(' | ' + msg) + L.append(" | " + msg) - msg = ''.join(L) + msg = "".join(L) sys.stdout.write(msg) # for i in range(term_width-int(TOTAL_BAR_LENGTH)-len(msg)-3): # sys.stdout.write(' ') @@ -122,76 +134,74 @@ def progress_bar(current, total, msg=None): # Go back to the center of the bar. # for i in range(term_width-int(TOTAL_BAR_LENGTH/2)+2): # sys.stdout.write('\b') - sys.stdout.write(' %d/%d ' % (current+1, total)) + sys.stdout.write(" %d/%d " % (current + 1, total)) - if current < total-1: - sys.stdout.write('\r') + if current < total - 1: + sys.stdout.write("\r") else: - sys.stdout.write('\n') + sys.stdout.write("\n") sys.stdout.flush() def format_time(seconds): - days = int(seconds / 3600/24) - seconds = seconds - days*3600*24 + days = int(seconds / 3600 / 24) + seconds = seconds - days * 3600 * 24 hours = int(seconds / 3600) - seconds = seconds - hours*3600 + seconds = seconds - hours * 3600 minutes = int(seconds / 60) - seconds = seconds - minutes*60 + seconds = seconds - minutes * 60 secondsf = int(seconds) seconds = seconds - secondsf - millis = int(seconds*1000) + millis = int(seconds * 1000) - f = '' + f = "" i = 1 if days > 0: - f += str(days) + 'D' + f += str(days) + "D" i += 1 if hours > 0 and i <= 2: - f += str(hours) + 'h' + f += str(hours) + "h" i += 1 if minutes > 0 and i <= 2: - f += str(minutes) + 'm' + f += str(minutes) + "m" i += 1 if secondsf > 0 and i <= 2: - f += str(secondsf) + 's' + f += str(secondsf) + "s" i += 1 if millis > 0 and i <= 2: - f += str(millis) + 'ms' + f += str(millis) + "ms" i += 1 - if f == '': - f = '0ms' + if f == "": + f = "0ms" return f def save_model(net, epoch, path, acc, is_best, **kwargs): state = { - 'net': net.state_dict(), - 'epoch': epoch, - 'acc': acc + "net": net.state_dict(), + "epoch": epoch, + "acc": acc, } for key, value in kwargs.items(): state[key] = value filepath = os.path.join(path, "last_checkpoint.pth") torch.save(state, filepath) if is_best: - shutil.copyfile(filepath, os.path.join(path, 'best_checkpoint.pth')) - + shutil.copyfile(filepath, os.path.join(path, "best_checkpoint.pth")) def save_args(args): - file = open(os.path.join(args.checkpoint, 'args.txt'), "w") + file = open(os.path.join(args.checkpoint, "args.txt"), "w") for k, v in vars(args).items(): file.write(f"{k}:\t {v}\n") file.close() - def set_seed(seed=None): if seed is None: return random.seed(seed) - os.environ['PYTHONHASHSEED'] = ("%s" % seed) + os.environ["PYTHONHASHSEED"] = "%s" % seed np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed(seed) @@ -200,15 +210,14 @@ def set_seed(seed=None): torch.backends.cudnn.deterministic = True - # create a file and write the text into it -class IOStream(): +class IOStream: def __init__(self, path): - self.f = open(path, 'a') + self.f = open(path, "a") def cprint(self, text): print(text) - self.f.write(text+'\n') + self.f.write(text + "\n") self.f.flush() def close(self): @@ -216,8 +225,7 @@ class IOStream(): def cal_loss(pred, gold, smoothing=True): - ''' Calculate cross entropy loss, apply label smoothing if needed. ''' - + """Calculate cross entropy loss, apply label smoothing if needed.""" gold = gold.contiguous().view(-1) if smoothing: @@ -230,6 +238,6 @@ def cal_loss(pred, gold, smoothing=True): loss = -(one_hot * log_prb).sum(dim=1).mean() else: - loss = F.cross_entropy(pred, gold, reduction='mean') + loss = F.cross_entropy(pred, gold, reduction="mean") return loss diff --git a/classification_ModelNet40/utils/progress/progress/__init__.py b/classification_ModelNet40/utils/progress/progress/__init__.py index 09dfc1e..1166f4b 100644 --- a/classification_ModelNet40/utils/progress/progress/__init__.py +++ b/classification_ModelNet40/utils/progress/progress/__init__.py @@ -12,7 +12,6 @@ # ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF # OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE. -from __future__ import division from collections import deque from datetime import timedelta @@ -20,13 +19,12 @@ from math import ceil from sys import stderr from time import time - -__version__ = '1.3' +__version__ = "1.3" -class Infinite(object): +class Infinite: file = stderr - sma_window = 10 # Simple Moving Average window + sma_window = 10 # Simple Moving Average window def __init__(self, *args, **kwargs): self.index = 0 @@ -38,7 +36,7 @@ class Infinite(object): setattr(self, key, val) def __getitem__(self, key): - if key.startswith('_'): + if key.startswith("_"): return None return getattr(self, key, None) @@ -83,8 +81,8 @@ class Infinite(object): class Progress(Infinite): def __init__(self, *args, **kwargs): - super(Progress, self).__init__(*args, **kwargs) - self.max = kwargs.get('max', 100) + super().__init__(*args, **kwargs) + self.max = kwargs.get("max", 100) @property def eta(self): diff --git a/classification_ModelNet40/utils/progress/progress/bar.py b/classification_ModelNet40/utils/progress/progress/bar.py index 5ee968f..93ec858 100644 --- a/classification_ModelNet40/utils/progress/progress/bar.py +++ b/classification_ModelNet40/utils/progress/progress/bar.py @@ -1,5 +1,3 @@ -# -*- coding: utf-8 -*- - # Copyright (c) 2012 Giorgos Verigakis # # Permission to use, copy, modify, and distribute this software for any @@ -14,19 +12,18 @@ # ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF # OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE. -from __future__ import unicode_literals from . import Progress from .helpers import WritelnMixin class Bar(WritelnMixin, Progress): width = 32 - message = '' - suffix = '%(index)d/%(max)d' - bar_prefix = ' |' - bar_suffix = '| ' - empty_fill = ' ' - fill = '#' + message = "" + suffix = "%(index)d/%(max)d" + bar_prefix = " |" + bar_suffix = "| " + empty_fill = " " + fill = "#" hide_cursor = True def update(self): @@ -37,52 +34,50 @@ class Bar(WritelnMixin, Progress): bar = self.fill * filled_length empty = self.empty_fill * empty_length suffix = self.suffix % self - line = ''.join([message, self.bar_prefix, bar, empty, self.bar_suffix, - suffix]) + line = "".join([message, self.bar_prefix, bar, empty, self.bar_suffix, suffix]) self.writeln(line) class ChargingBar(Bar): - suffix = '%(percent)d%%' - bar_prefix = ' ' - bar_suffix = ' ' - empty_fill = '∙' - fill = '█' + suffix = "%(percent)d%%" + bar_prefix = " " + bar_suffix = " " + empty_fill = "∙" + fill = "█" class FillingSquaresBar(ChargingBar): - empty_fill = '▢' - fill = '▣' + empty_fill = "▢" + fill = "▣" class FillingCirclesBar(ChargingBar): - empty_fill = '◯' - fill = '◉' + empty_fill = "◯" + fill = "◉" class IncrementalBar(Bar): - phases = (' ', '▏', '▎', '▍', '▌', '▋', '▊', '▉', '█') + phases = (" ", "▏", "▎", "▍", "▌", "▋", "▊", "▉", "█") def update(self): nphases = len(self.phases) filled_len = self.width * self.progress - nfull = int(filled_len) # Number of full chars + nfull = int(filled_len) # Number of full chars phase = int((filled_len - nfull) * nphases) # Phase of last char - nempty = self.width - nfull # Number of empty chars + nempty = self.width - nfull # Number of empty chars message = self.message % self bar = self.phases[-1] * nfull - current = self.phases[phase] if phase > 0 else '' + current = self.phases[phase] if phase > 0 else "" empty = self.empty_fill * max(0, nempty - len(current)) suffix = self.suffix % self - line = ''.join([message, self.bar_prefix, bar, current, empty, - self.bar_suffix, suffix]) + line = "".join([message, self.bar_prefix, bar, current, empty, self.bar_suffix, suffix]) self.writeln(line) class PixelBar(IncrementalBar): - phases = ('⡀', '⡄', '⡆', '⡇', '⣇', '⣧', '⣷', '⣿') + phases = ("⡀", "⡄", "⡆", "⡇", "⣇", "⣧", "⣷", "⣿") class ShadyBar(IncrementalBar): - phases = (' ', '░', '▒', '▓', '█') + phases = (" ", "░", "▒", "▓", "█") diff --git a/classification_ModelNet40/utils/progress/progress/counter.py b/classification_ModelNet40/utils/progress/progress/counter.py index 6b45a1e..90da4ff 100644 --- a/classification_ModelNet40/utils/progress/progress/counter.py +++ b/classification_ModelNet40/utils/progress/progress/counter.py @@ -1,5 +1,3 @@ -# -*- coding: utf-8 -*- - # Copyright (c) 2012 Giorgos Verigakis # # Permission to use, copy, modify, and distribute this software for any @@ -14,13 +12,12 @@ # ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF # OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE. -from __future__ import unicode_literals from . import Infinite, Progress from .helpers import WriteMixin class Counter(WriteMixin, Infinite): - message = '' + message = "" hide_cursor = True def update(self): @@ -35,7 +32,7 @@ class Countdown(WriteMixin, Progress): class Stack(WriteMixin, Progress): - phases = (' ', '▁', '▂', '▃', '▄', '▅', '▆', '▇', '█') + phases = (" ", "▁", "▂", "▃", "▄", "▅", "▆", "▇", "█") hide_cursor = True def update(self): @@ -45,4 +42,4 @@ class Stack(WriteMixin, Progress): class Pie(Stack): - phases = ('○', '◔', '◑', '◕', '●') + phases = ("○", "◔", "◑", "◕", "●") diff --git a/classification_ModelNet40/utils/progress/progress/helpers.py b/classification_ModelNet40/utils/progress/progress/helpers.py index 9ed90b2..04abe24 100644 --- a/classification_ModelNet40/utils/progress/progress/helpers.py +++ b/classification_ModelNet40/utils/progress/progress/helpers.py @@ -12,78 +12,76 @@ # ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF # OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE. -from __future__ import print_function + +HIDE_CURSOR = "\x1b[?25l" +SHOW_CURSOR = "\x1b[?25h" -HIDE_CURSOR = '\x1b[?25l' -SHOW_CURSOR = '\x1b[?25h' - - -class WriteMixin(object): +class WriteMixin: hide_cursor = False def __init__(self, message=None, **kwargs): - super(WriteMixin, self).__init__(**kwargs) + super().__init__(**kwargs) self._width = 0 if message: self.message = message if self.file.isatty(): if self.hide_cursor: - print(HIDE_CURSOR, end='', file=self.file) - print(self.message, end='', file=self.file) + print(HIDE_CURSOR, end="", file=self.file) + print(self.message, end="", file=self.file) self.file.flush() def write(self, s): if self.file.isatty(): - b = '\b' * self._width + b = "\b" * self._width c = s.ljust(self._width) - print(b + c, end='', file=self.file) + print(b + c, end="", file=self.file) self._width = max(self._width, len(s)) self.file.flush() def finish(self): if self.file.isatty() and self.hide_cursor: - print(SHOW_CURSOR, end='', file=self.file) + print(SHOW_CURSOR, end="", file=self.file) -class WritelnMixin(object): +class WritelnMixin: hide_cursor = False def __init__(self, message=None, **kwargs): - super(WritelnMixin, self).__init__(**kwargs) + super().__init__(**kwargs) if message: self.message = message if self.file.isatty() and self.hide_cursor: - print(HIDE_CURSOR, end='', file=self.file) + print(HIDE_CURSOR, end="", file=self.file) def clearln(self): if self.file.isatty(): - print('\r\x1b[K', end='', file=self.file) + print("\r\x1b[K", end="", file=self.file) def writeln(self, line): if self.file.isatty(): self.clearln() - print(line, end='', file=self.file) + print(line, end="", file=self.file) self.file.flush() def finish(self): if self.file.isatty(): print(file=self.file) if self.hide_cursor: - print(SHOW_CURSOR, end='', file=self.file) + print(SHOW_CURSOR, end="", file=self.file) -from signal import signal, SIGINT +from signal import SIGINT, signal from sys import exit -class SigIntMixin(object): - """Registers a signal handler that calls finish on SIGINT""" +class SigIntMixin: + """Registers a signal handler that calls finish on SIGINT.""" def __init__(self, *args, **kwargs): - super(SigIntMixin, self).__init__(*args, **kwargs) + super().__init__(*args, **kwargs) signal(SIGINT, self._sigint_handler) def _sigint_handler(self, signum, frame): diff --git a/classification_ModelNet40/utils/progress/progress/spinner.py b/classification_ModelNet40/utils/progress/progress/spinner.py index 464c7b2..77d0f49 100644 --- a/classification_ModelNet40/utils/progress/progress/spinner.py +++ b/classification_ModelNet40/utils/progress/progress/spinner.py @@ -1,5 +1,3 @@ -# -*- coding: utf-8 -*- - # Copyright (c) 2012 Giorgos Verigakis # # Permission to use, copy, modify, and distribute this software for any @@ -14,14 +12,13 @@ # ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF # OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE. -from __future__ import unicode_literals from . import Infinite from .helpers import WriteMixin class Spinner(WriteMixin, Infinite): - message = '' - phases = ('-', '\\', '|', '/') + message = "" + phases = ("-", "\\", "|", "/") hide_cursor = True def update(self): @@ -30,15 +27,16 @@ class Spinner(WriteMixin, Infinite): class PieSpinner(Spinner): - phases = ['◷', '◶', '◵', '◴'] + phases = ["◷", "◶", "◵", "◴"] class MoonSpinner(Spinner): - phases = ['◑', '◒', '◐', '◓'] + phases = ["◑", "◒", "◐", "◓"] class LineSpinner(Spinner): - phases = ['⎺', '⎻', '⎼', '⎽', '⎼', '⎻'] + phases = ["⎺", "⎻", "⎼", "⎽", "⎼", "⎻"] + class PixelSpinner(Spinner): - phases = ['⣾','⣷', '⣯', '⣟', '⡿', '⢿', '⣻', '⣽'] + phases = ["⣾", "⣷", "⣯", "⣟", "⡿", "⢿", "⣻", "⣽"] diff --git a/classification_ModelNet40/utils/progress/setup.py b/classification_ModelNet40/utils/progress/setup.py index c877781..4d935fd 100755 --- a/classification_ModelNet40/utils/progress/setup.py +++ b/classification_ModelNet40/utils/progress/setup.py @@ -1,29 +1,27 @@ #!/usr/bin/env python +import progress from setuptools import setup -import progress - - setup( - name='progress', + name="progress", version=progress.__version__, - description='Easy to use progress bars', - long_description=open('README.rst').read(), - author='Giorgos Verigakis', - author_email='verigak@gmail.com', - url='http://github.com/verigak/progress/', - license='ISC', - packages=['progress'], + description="Easy to use progress bars", + long_description=open("README.rst").read(), + author="Giorgos Verigakis", + author_email="verigak@gmail.com", + url="http://github.com/verigak/progress/", + license="ISC", + packages=["progress"], classifiers=[ - 'Environment :: Console', - 'Intended Audience :: Developers', - 'License :: OSI Approved :: ISC License (ISCL)', - 'Programming Language :: Python :: 2.6', - 'Programming Language :: Python :: 2.7', - 'Programming Language :: Python :: 3.3', - 'Programming Language :: Python :: 3.4', - 'Programming Language :: Python :: 3.5', - 'Programming Language :: Python :: 3.6', - ] + "Environment :: Console", + "Intended Audience :: Developers", + "License :: OSI Approved :: ISC License (ISCL)", + "Programming Language :: Python :: 2.6", + "Programming Language :: Python :: 2.7", + "Programming Language :: Python :: 3.3", + "Programming Language :: Python :: 3.4", + "Programming Language :: Python :: 3.5", + "Programming Language :: Python :: 3.6", + ], ) diff --git a/classification_ModelNet40/utils/progress/test_progress.py b/classification_ModelNet40/utils/progress/test_progress.py index 0f68b01..1f200c3 100755 --- a/classification_ModelNet40/utils/progress/test_progress.py +++ b/classification_ModelNet40/utils/progress/test_progress.py @@ -1,16 +1,12 @@ #!/usr/bin/env python -from __future__ import print_function import random import time -from progress.bar import (Bar, ChargingBar, FillingSquaresBar, - FillingCirclesBar, IncrementalBar, PixelBar, - ShadyBar) -from progress.spinner import (Spinner, PieSpinner, MoonSpinner, LineSpinner, - PixelSpinner) -from progress.counter import Counter, Countdown, Stack, Pie +from progress.bar import Bar, ChargingBar, FillingCirclesBar, FillingSquaresBar, IncrementalBar, PixelBar, ShadyBar +from progress.counter import Countdown, Counter, Pie, Stack +from progress.spinner import LineSpinner, MoonSpinner, PieSpinner, PixelSpinner, Spinner def sleep(): @@ -20,29 +16,29 @@ def sleep(): for bar_cls in (Bar, ChargingBar, FillingSquaresBar, FillingCirclesBar): - suffix = '%(index)d/%(max)d [%(elapsed)d / %(eta)d / %(eta_td)s]' + suffix = "%(index)d/%(max)d [%(elapsed)d / %(eta)d / %(eta_td)s]" bar = bar_cls(bar_cls.__name__, suffix=suffix) - for i in bar.iter(range(200)): + for _i in bar.iter(range(200)): sleep() for bar_cls in (IncrementalBar, PixelBar, ShadyBar): - suffix = '%(percent)d%% [%(elapsed_td)s / %(eta)d / %(eta_td)s]' + suffix = "%(percent)d%% [%(elapsed_td)s / %(eta)d / %(eta_td)s]" bar = bar_cls(bar_cls.__name__, suffix=suffix) - for i in bar.iter(range(200)): + for _i in bar.iter(range(200)): sleep() for spin in (Spinner, PieSpinner, MoonSpinner, LineSpinner, PixelSpinner): - for i in spin(spin.__name__ + ' ').iter(range(100)): + for _i in spin(spin.__name__ + " ").iter(range(100)): sleep() print() for singleton in (Counter, Countdown, Stack, Pie): - for i in singleton(singleton.__name__ + ' ').iter(range(100)): + for _i in singleton(singleton.__name__ + " ").iter(range(100)): sleep() print() -bar = IncrementalBar('Random', suffix='%(index)d') -for i in range(100): +bar = IncrementalBar("Random", suffix="%(index)d") +for _i in range(100): bar.goto(random.randint(0, 100)) sleep() bar.finish() diff --git a/classification_ModelNet40/voting.py b/classification_ModelNet40/voting.py index 9eda038..330fcfd 100644 --- a/classification_ModelNet40/voting.py +++ b/classification_ModelNet40/voting.py @@ -1,47 +1,52 @@ import argparse -import os import datetime +import os + +import models as models +import numpy as np +import sklearn.metrics as metrics import torch -import torch.nn.parallel import torch.backends.cudnn as cudnn +import torch.nn.functional as F +import torch.nn.parallel import torch.optim import torch.utils.data import torch.utils.data.distributed -from torch.utils.data import DataLoader -import models as models -from utils import progress_bar, IOStream from data import ModelNet40 -import sklearn.metrics as metrics from helper import cal_loss -import numpy as np -import torch.nn.functional as F +from torch.utils.data import DataLoader +from utils import IOStream, progress_bar -model_names = sorted(name for name in models.__dict__ - if callable(models.__dict__[name])) +model_names = sorted(name for name in models.__dict__ if callable(models.__dict__[name])) def parse_args(): """Parameters""" - parser = argparse.ArgumentParser('training') - parser.add_argument('-c', '--checkpoint', type=str, metavar='PATH', - help='path to save checkpoint (default: checkpoint)') - parser.add_argument('--msg', type=str, help='message after checkpoint') - parser.add_argument('--batch_size', type=int, default=32, help='batch size in training') - parser.add_argument('--model', default='model31A', help='model name [default: pointnet_cls]') - parser.add_argument('--num_classes', default=40, type=int, choices=[10, 40], help='training on ModelNet10/40') - parser.add_argument('--num_points', type=int, default=1024, help='Point Number') - parser.add_argument('--seed', type=int, help='random seed (default: 1)') + parser = argparse.ArgumentParser("training") + parser.add_argument( + "-c", + "--checkpoint", + type=str, + metavar="PATH", + help="path to save checkpoint (default: checkpoint)", + ) + parser.add_argument("--msg", type=str, help="message after checkpoint") + parser.add_argument("--batch_size", type=int, default=32, help="batch size in training") + parser.add_argument("--model", default="model31A", help="model name [default: pointnet_cls]") + parser.add_argument("--num_classes", default=40, type=int, choices=[10, 40], help="training on ModelNet10/40") + parser.add_argument("--num_points", type=int, default=1024, help="Point Number") + parser.add_argument("--seed", type=int, help="random seed (default: 1)") # Voting evaluation, referring: https://github.com/CVMI-Lab/PAConv/blob/main/obj_cls/eval_voting.py - parser.add_argument('--NUM_PEPEAT', type=int, default=300) - parser.add_argument('--NUM_VOTE', type=int, default=10) + parser.add_argument("--NUM_PEPEAT", type=int, default=300) + parser.add_argument("--NUM_VOTE", type=int, default=10) - parser.add_argument('--validate', action='store_true', help='Validate the original testing result.') + parser.add_argument("--validate", action="store_true", help="Validate the original testing result.") return parser.parse_args() -class PointcloudScale(object): # input random scaling - def __init__(self, scale_low=2. / 3., scale_high=3. / 2.): +class PointcloudScale: # input random scaling + def __init__(self, scale_low=2.0 / 3.0, scale_high=3.0 / 2.0): self.scale_low = scale_low self.scale_high = scale_high @@ -68,45 +73,52 @@ def main(): torch.set_printoptions(10) torch.backends.cudnn.benchmark = False torch.backends.cudnn.deterministic = True - os.environ['PYTHONHASHSEED'] = str(args.seed) + os.environ["PYTHONHASHSEED"] = str(args.seed) if torch.cuda.is_available(): - device = 'cuda' + device = "cuda" else: - device = 'cpu' + device = "cpu" print(f"==> Using device: {device}") if args.msg is None: - message = str(datetime.datetime.now().strftime('-%Y%m%d%H%M%S')) + message = str(datetime.datetime.now().strftime("-%Y%m%d%H%M%S")) else: message = "-" + args.msg - args.checkpoint = 'checkpoints/' + args.model + message + args.checkpoint = "checkpoints/" + args.model + message - print('==> Preparing data..') - test_loader = DataLoader(ModelNet40(partition='test', num_points=args.num_points), num_workers=4, - batch_size=args.batch_size // 2, shuffle=False, drop_last=False) + print("==> Preparing data..") + test_loader = DataLoader( + ModelNet40(partition="test", num_points=args.num_points), + num_workers=4, + batch_size=args.batch_size // 2, + shuffle=False, + drop_last=False, + ) # Model - print('==> Building model..') + print("==> Building model..") net = models.__dict__[args.model]() criterion = cal_loss net = net.to(device) - checkpoint_path = os.path.join(args.checkpoint, 'best_checkpoint.pth') - checkpoint = torch.load(checkpoint_path, map_location=torch.device('cpu')) + checkpoint_path = os.path.join(args.checkpoint, "best_checkpoint.pth") + checkpoint = torch.load(checkpoint_path, map_location=torch.device("cpu")) # criterion = criterion.to(device) - if device == 'cuda': + if device == "cuda": net = torch.nn.DataParallel(net) cudnn.benchmark = True - net.load_state_dict(checkpoint['net']) + net.load_state_dict(checkpoint["net"]) if args.validate: test_out = validate(net, test_loader, criterion, device) print(f"Vanilla out: {test_out}") - print(f"Note 1: Please also load the random seed parameter (if forgot, see out.txt).\n" - f"Note 2: This result may vary little on different GPUs (and number of GPUs), we tested 2080Ti, P100, and V100.\n" - f"[note : Original result is achieved with V100 GPUs.]\n\n\n") + print( + "Note 1: Please also load the random seed parameter (if forgot, see out.txt).\n" + "Note 2: This result may vary little on different GPUs (and number of GPUs), we tested 2080Ti, P100, and V100.\n" + "[note : Original result is achieved with V100 GPUs.]\n\n\n", + ) # Interestingly, we get original best_test_acc on 4 V100 gpus, but this model is trained on one V100 gpu. # On different GPUs, and different number of GPUs, both OA and mean_acc vary a little. # Also, the batch size also affect the testing results, could not understand. - print(f"===> start voting evaluation...") + print("===> start voting evaluation...") voting(net, test_loader, device, args) @@ -130,23 +142,28 @@ def validate(net, testloader, criterion, device): test_pred.append(preds.detach().cpu().numpy()) total += label.size(0) correct += preds.eq(label).sum().item() - progress_bar(batch_idx, len(testloader), 'Loss: %.3f | Acc: %.3f%% (%d/%d)' - % (test_loss / (batch_idx + 1), 100. * correct / total, correct, total)) + progress_bar( + batch_idx, + len(testloader), + "Loss: %.3f | Acc: %.3f%% (%d/%d)" + % (test_loss / (batch_idx + 1), 100.0 * correct / total, correct, total), + ) time_cost = int((datetime.datetime.now() - time_cost).total_seconds()) test_true = np.concatenate(test_true) test_pred = np.concatenate(test_pred) return { "loss": float("%.3f" % (test_loss / (batch_idx + 1))), - "acc": float("%.3f" % (100. * metrics.accuracy_score(test_true, test_pred))), - "acc_avg": float("%.3f" % (100. * metrics.balanced_accuracy_score(test_true, test_pred))), - "time": time_cost + "acc": float("%.3f" % (100.0 * metrics.accuracy_score(test_true, test_pred))), + "acc_avg": float("%.3f" % (100.0 * metrics.balanced_accuracy_score(test_true, test_pred))), + "time": time_cost, } def voting(net, testloader, device, args): - name = '/evaluate_voting' + str(datetime.datetime.now().strftime('-%Y%m%d%H%M%S')) + 'seed_' + str( - args.seed) + '.log' + name = ( + "/evaluate_voting" + str(datetime.datetime.now().strftime("-%Y%m%d%H%M%S")) + "seed_" + str(args.seed) + ".log" + ) io = IOStream(args.checkpoint + name) io.cprint(str(args)) @@ -161,7 +178,7 @@ def voting(net, testloader, device, args): test_true = [] test_pred = [] - for batch_idx, (data, label) in enumerate(testloader): + for _batch_idx, (data, label) in enumerate(testloader): data, label = data.to(device), label.to(device).squeeze() pred = 0 for v in range(args.NUM_VOTE): @@ -178,19 +195,24 @@ def voting(net, testloader, device, args): test_pred.append(pred_choice.detach().cpu().numpy()) test_true = np.concatenate(test_true) test_pred = np.concatenate(test_pred) - test_acc = 100. * metrics.accuracy_score(test_true, test_pred) - test_mean_acc = 100. * metrics.balanced_accuracy_score(test_true, test_pred) + test_acc = 100.0 * metrics.accuracy_score(test_true, test_pred) + test_mean_acc = 100.0 * metrics.balanced_accuracy_score(test_true, test_pred) if test_acc > best_acc: best_acc = test_acc if test_mean_acc > best_mean_acc: best_mean_acc = test_mean_acc - outstr = 'Voting %d, test acc: %.3f, test mean acc: %.3f, [current best(all_acc: %.3f mean_acc: %.3f)]' % \ - (i, test_acc, test_mean_acc, best_acc, best_mean_acc) + outstr = "Voting %d, test acc: %.3f, test mean acc: %.3f, [current best(all_acc: %.3f mean_acc: %.3f)]" % ( + i, + test_acc, + test_mean_acc, + best_acc, + best_mean_acc, + ) io.cprint(outstr) - final_outstr = 'Final voting test acc: %.6f,' % (best_acc * 100) + final_outstr = "Final voting test acc: %.6f," % (best_acc * 100) io.cprint(final_outstr) -if __name__ == '__main__': +if __name__ == "__main__": main() diff --git a/classification_ScanObjectNN/ScanObjectNN.py b/classification_ScanObjectNN/ScanObjectNN.py index b479141..cca8258 100644 --- a/classification_ScanObjectNN/ScanObjectNN.py +++ b/classification_ScanObjectNN/ScanObjectNN.py @@ -1,10 +1,7 @@ -""" -ScanObjectNN download: http://103.24.77.34/scanobjectnn/h5_files.zip -""" +"""ScanObjectNN download: http://103.24.77.34/scanobjectnn/h5_files.zip.""" import os -import sys -import glob + import h5py import numpy as np from torch.utils.data import Dataset @@ -14,18 +11,18 @@ os.environ["HDF5_USE_FILE_LOCKING"] = "FALSE" def download(): BASE_DIR = os.path.dirname(os.path.abspath(__file__)) - DATA_DIR = os.path.join(BASE_DIR, 'data') + DATA_DIR = os.path.join(BASE_DIR, "data") if not os.path.exists(DATA_DIR): os.mkdir(DATA_DIR) - if not os.path.exists(os.path.join(DATA_DIR, 'h5_files')): + if not os.path.exists(os.path.join(DATA_DIR, "h5_files")): # note that this link only contains the hardest perturbed variant (PB_T50_RS). # for full versions, consider the following link. - www = 'https://web.northeastern.edu/smilelab/xuma/datasets/h5_files.zip' + www = "https://web.northeastern.edu/smilelab/xuma/datasets/h5_files.zip" # www = 'http://103.24.77.34/scanobjectnn/h5_files.zip' zipfile = os.path.basename(www) - os.system('wget %s --no-check-certificate; unzip %s' % (www, zipfile)) - os.system('mv %s %s' % (zipfile[:-4], DATA_DIR)) - os.system('rm %s' % (zipfile)) + os.system(f"wget {www} --no-check-certificate; unzip {zipfile}") + os.system(f"mv {zipfile[:-4]} {DATA_DIR}") + os.system("rm %s" % (zipfile)) def load_scanobjectnn_data(partition): @@ -34,10 +31,10 @@ def load_scanobjectnn_data(partition): all_data = [] all_label = [] - h5_name = BASE_DIR + '/data/h5_files/main_split/' + partition + '_objectdataset_augmentedrot_scale75.h5' + h5_name = BASE_DIR + "/data/h5_files/main_split/" + partition + "_objectdataset_augmentedrot_scale75.h5" f = h5py.File(h5_name, mode="r") - data = f['data'][:].astype('float32') - label = f['label'][:].astype('int64') + data = f["data"][:].astype("float32") + label = f["label"][:].astype("int64") f.close() all_data.append(data) all_label.append(label) @@ -47,23 +44,22 @@ def load_scanobjectnn_data(partition): def translate_pointcloud(pointcloud): - xyz1 = np.random.uniform(low=2. / 3., high=3. / 2., size=[3]) + xyz1 = np.random.uniform(low=2.0 / 3.0, high=3.0 / 2.0, size=[3]) xyz2 = np.random.uniform(low=-0.2, high=0.2, size=[3]) - translated_pointcloud = np.add(np.multiply(pointcloud, xyz1), xyz2).astype('float32') - return translated_pointcloud + return np.add(np.multiply(pointcloud, xyz1), xyz2).astype("float32") class ScanObjectNN(Dataset): - def __init__(self, num_points, partition='training'): + def __init__(self, num_points, partition="training"): self.data, self.label = load_scanobjectnn_data(partition) self.num_points = num_points self.partition = partition def __getitem__(self, item): - pointcloud = self.data[item][:self.num_points] + pointcloud = self.data[item][: self.num_points] label = self.label[item] - if self.partition == 'training': + if self.partition == "training": pointcloud = translate_pointcloud(pointcloud) np.random.shuffle(pointcloud) return pointcloud, label @@ -72,9 +68,9 @@ class ScanObjectNN(Dataset): return self.data.shape[0] -if __name__ == '__main__': +if __name__ == "__main__": train = ScanObjectNN(1024) - test = ScanObjectNN(1024, 'test') + test = ScanObjectNN(1024, "test") for data, label in train: print(data.shape) print(label) diff --git a/classification_ScanObjectNN/main.py b/classification_ScanObjectNN/main.py index 60a4156..3d989ae 100644 --- a/classification_ScanObjectNN/main.py +++ b/classification_ScanObjectNN/main.py @@ -1,45 +1,50 @@ -""" -for training with resume functions. +"""for training with resume functions. Usage: python main.py --model PointNet --msg demo or -CUDA_VISIBLE_DEVICES=0 nohup python main.py --model PointNet --msg demo > nohup/PointNet_demo.out & +CUDA_VISIBLE_DEVICES=0 nohup python main.py --model PointNet --msg demo > nohup/PointNet_demo.out &. """ import argparse -import os -import logging import datetime +import logging +import os + +import models as models +import numpy as np +import sklearn.metrics as metrics import torch -import torch.nn.parallel import torch.backends.cudnn as cudnn +import torch.nn.parallel import torch.optim import torch.utils.data import torch.utils.data.distributed -from torch.utils.data import DataLoader -import models as models -from utils import Logger, mkdir_p, progress_bar, save_model, save_args, cal_loss from ScanObjectNN import ScanObjectNN from torch.optim.lr_scheduler import CosineAnnealingLR -import sklearn.metrics as metrics -import numpy as np +from torch.utils.data import DataLoader +from utils import Logger, cal_loss, mkdir_p, progress_bar, save_args, save_model def parse_args(): """Parameters""" - parser = argparse.ArgumentParser('training') - parser.add_argument('-c', '--checkpoint', type=str, metavar='PATH', - help='path to save checkpoint (default: checkpoint)') - parser.add_argument('--msg', type=str, help='message after checkpoint') - parser.add_argument('--batch_size', type=int, default=32, help='batch size in training') - parser.add_argument('--model', default='PointNet', help='model name [default: pointnet_cls]') - parser.add_argument('--num_classes', default=15, type=int, help='default value for classes of ScanObjectNN') - parser.add_argument('--epoch', default=200, type=int, help='number of epoch in training') - parser.add_argument('--num_points', type=int, default=1024, help='Point Number') - parser.add_argument('--learning_rate', default=0.01, type=float, help='learning rate in training') - parser.add_argument('--weight_decay', type=float, default=1e-4, help='decay rate') - parser.add_argument('--smoothing', action='store_true', default=False, help='loss smoothing') - parser.add_argument('--seed', type=int, help='random seed') - parser.add_argument('--workers', default=4, type=int, help='workers') + parser = argparse.ArgumentParser("training") + parser.add_argument( + "-c", + "--checkpoint", + type=str, + metavar="PATH", + help="path to save checkpoint (default: checkpoint)", + ) + parser.add_argument("--msg", type=str, help="message after checkpoint") + parser.add_argument("--batch_size", type=int, default=32, help="batch size in training") + parser.add_argument("--model", default="PointNet", help="model name [default: pointnet_cls]") + parser.add_argument("--num_classes", default=15, type=int, help="default value for classes of ScanObjectNN") + parser.add_argument("--epoch", default=200, type=int, help="number of epoch in training") + parser.add_argument("--num_points", type=int, default=1024, help="Point Number") + parser.add_argument("--learning_rate", default=0.01, type=float, help="learning rate in training") + parser.add_argument("--weight_decay", type=float, default=1e-4, help="decay rate") + parser.add_argument("--smoothing", action="store_true", default=False, help="loss smoothing") + parser.add_argument("--seed", type=int, help="random seed") + parser.add_argument("--workers", default=4, type=int, help="workers") return parser.parse_args() @@ -49,23 +54,23 @@ def main(): if args.seed is not None: torch.manual_seed(args.seed) if torch.cuda.is_available(): - device = 'cuda' + device = "cuda" if args.seed is not None: torch.cuda.manual_seed(args.seed) else: - device = 'cpu' - time_str = str(datetime.datetime.now().strftime('-%Y%m%d%H%M%S')) + device = "cpu" + time_str = str(datetime.datetime.now().strftime("-%Y%m%d%H%M%S")) if args.msg is None: message = time_str else: message = "-" + args.msg - args.checkpoint = 'checkpoints/' + args.model + message + args.checkpoint = "checkpoints/" + args.model + message if not os.path.isdir(args.checkpoint): mkdir_p(args.checkpoint) screen_logger = logging.getLogger("Model") screen_logger.setLevel(logging.INFO) - formatter = logging.Formatter('%(message)s') + formatter = logging.Formatter("%(message)s") file_handler = logging.FileHandler(os.path.join(args.checkpoint, "out.txt")) file_handler.setLevel(logging.INFO) file_handler.setFormatter(formatter) @@ -77,19 +82,19 @@ def main(): # Model printf(f"args: {args}") - printf('==> Building model..') + printf("==> Building model..") net = models.__dict__[args.model](num_classes=args.num_classes) criterion = cal_loss net = net.to(device) # criterion = criterion.to(device) - if device == 'cuda': + if device == "cuda": net = torch.nn.DataParallel(net) cudnn.benchmark = True - best_test_acc = 0. # best test accuracy - best_train_acc = 0. - best_test_acc_avg = 0. - best_train_acc_avg = 0. + best_test_acc = 0.0 # best test accuracy + best_train_acc = 0.0 + best_test_acc_avg = 0.0 + best_train_acc_avg = 0.0 best_test_loss = float("inf") best_train_loss = float("inf") start_epoch = 0 # start from epoch 0 or last checkpoint epoch @@ -97,30 +102,49 @@ def main(): if not os.path.isfile(os.path.join(args.checkpoint, "last_checkpoint.pth")): save_args(args) - logger = Logger(os.path.join(args.checkpoint, 'log.txt'), title="ModelNet" + args.model) - logger.set_names(["Epoch-Num", 'Learning-Rate', - 'Train-Loss', 'Train-acc-B', 'Train-acc', - 'Valid-Loss', 'Valid-acc-B', 'Valid-acc']) + logger = Logger(os.path.join(args.checkpoint, "log.txt"), title="ModelNet" + args.model) + logger.set_names( + [ + "Epoch-Num", + "Learning-Rate", + "Train-Loss", + "Train-acc-B", + "Train-acc", + "Valid-Loss", + "Valid-acc-B", + "Valid-acc", + ], + ) else: printf(f"Resuming last checkpoint from {args.checkpoint}") checkpoint_path = os.path.join(args.checkpoint, "last_checkpoint.pth") checkpoint = torch.load(checkpoint_path) - net.load_state_dict(checkpoint['net']) - start_epoch = checkpoint['epoch'] - best_test_acc = checkpoint['best_test_acc'] - best_train_acc = checkpoint['best_train_acc'] - best_test_acc_avg = checkpoint['best_test_acc_avg'] - best_train_acc_avg = checkpoint['best_train_acc_avg'] - best_test_loss = checkpoint['best_test_loss'] - best_train_loss = checkpoint['best_train_loss'] - logger = Logger(os.path.join(args.checkpoint, 'log.txt'), title="ModelNet" + args.model, resume=True) - optimizer_dict = checkpoint['optimizer'] + net.load_state_dict(checkpoint["net"]) + start_epoch = checkpoint["epoch"] + best_test_acc = checkpoint["best_test_acc"] + best_train_acc = checkpoint["best_train_acc"] + best_test_acc_avg = checkpoint["best_test_acc_avg"] + best_train_acc_avg = checkpoint["best_train_acc_avg"] + best_test_loss = checkpoint["best_test_loss"] + best_train_loss = checkpoint["best_train_loss"] + logger = Logger(os.path.join(args.checkpoint, "log.txt"), title="ModelNet" + args.model, resume=True) + optimizer_dict = checkpoint["optimizer"] - printf('==> Preparing data..') - train_loader = DataLoader(ScanObjectNN(partition='training', num_points=args.num_points), num_workers=args.workers, - batch_size=args.batch_size, shuffle=True, drop_last=True) - test_loader = DataLoader(ScanObjectNN(partition='test', num_points=args.num_points), num_workers=args.workers, - batch_size=args.batch_size, shuffle=True, drop_last=False) + printf("==> Preparing data..") + train_loader = DataLoader( + ScanObjectNN(partition="training", num_points=args.num_points), + num_workers=args.workers, + batch_size=args.batch_size, + shuffle=True, + drop_last=True, + ) + test_loader = DataLoader( + ScanObjectNN(partition="test", num_points=args.num_points), + num_workers=args.workers, + batch_size=args.batch_size, + shuffle=True, + drop_last=False, + ) optimizer = torch.optim.SGD(net.parameters(), lr=args.learning_rate, momentum=0.9, weight_decay=args.weight_decay) if optimizer_dict is not None: @@ -128,7 +152,7 @@ def main(): scheduler = CosineAnnealingLR(optimizer, args.epoch, eta_min=args.learning_rate / 100, last_epoch=start_epoch - 1) for epoch in range(start_epoch, args.epoch): - printf('Epoch(%d/%s) Learning Rate %s:' % (epoch + 1, args.epoch, optimizer.param_groups[0]['lr'])) + printf("Epoch(%d/%s) Learning Rate %s:" % (epoch + 1, args.epoch, optimizer.param_groups[0]["lr"])) train_out = train(net, train_loader, optimizer, criterion, device) # {"loss", "acc", "acc_avg", "time"} test_out = validate(net, test_loader, criterion, device) scheduler.step() @@ -147,31 +171,46 @@ def main(): best_train_loss = train_out["loss"] if (train_out["loss"] < best_train_loss) else best_train_loss save_model( - net, epoch, path=args.checkpoint, acc=test_out["acc"], is_best=is_best, + net, + epoch, + path=args.checkpoint, + acc=test_out["acc"], + is_best=is_best, best_test_acc=best_test_acc, # best test accuracy best_train_acc=best_train_acc, best_test_acc_avg=best_test_acc_avg, best_train_acc_avg=best_train_acc_avg, best_test_loss=best_test_loss, best_train_loss=best_train_loss, - optimizer=optimizer.state_dict() + optimizer=optimizer.state_dict(), + ) + logger.append( + [ + epoch, + optimizer.param_groups[0]["lr"], + train_out["loss"], + train_out["acc_avg"], + train_out["acc"], + test_out["loss"], + test_out["acc_avg"], + test_out["acc"], + ], ) - logger.append([epoch, optimizer.param_groups[0]['lr'], - train_out["loss"], train_out["acc_avg"], train_out["acc"], - test_out["loss"], test_out["acc_avg"], test_out["acc"]]) printf( - f"Training loss:{train_out['loss']} acc_avg:{train_out['acc_avg']}% acc:{train_out['acc']}% time:{train_out['time']}s") + f"Training loss:{train_out['loss']} acc_avg:{train_out['acc_avg']}% acc:{train_out['acc']}% time:{train_out['time']}s", + ) printf( f"Testing loss:{test_out['loss']} acc_avg:{test_out['acc_avg']}% " - f"acc:{test_out['acc']}% time:{test_out['time']}s [best test acc: {best_test_acc}%] \n\n") + f"acc:{test_out['acc']}% time:{test_out['time']}s [best test acc: {best_test_acc}%] \n\n", + ) logger.close() - printf(f"++++++++" * 2 + "Final results" + "++++++++" * 2) + printf("++++++++" * 2 + "Final results" + "++++++++" * 2) printf(f"++ Last Train time: {train_out['time']} | Last Test time: {test_out['time']} ++") printf(f"++ Best Train loss: {best_train_loss} | Best Test loss: {best_test_loss} ++") printf(f"++ Best Train acc_B: {best_train_acc_avg} | Best Test acc_B: {best_test_acc_avg} ++") printf(f"++ Best Train acc: {best_train_acc} | Best Test acc: {best_test_acc} ++") - printf(f"++++++++" * 5) + printf("++++++++" * 5) def train(net, trainloader, optimizer, criterion, device): @@ -199,17 +238,21 @@ def train(net, trainloader, optimizer, criterion, device): total += label.size(0) correct += preds.eq(label).sum().item() - progress_bar(batch_idx, len(trainloader), 'Loss: %.3f | Acc: %.3f%% (%d/%d)' - % (train_loss / (batch_idx + 1), 100. * correct / total, correct, total)) + progress_bar( + batch_idx, + len(trainloader), + "Loss: %.3f | Acc: %.3f%% (%d/%d)" + % (train_loss / (batch_idx + 1), 100.0 * correct / total, correct, total), + ) time_cost = int((datetime.datetime.now() - time_cost).total_seconds()) train_true = np.concatenate(train_true) train_pred = np.concatenate(train_pred) return { "loss": float("%.3f" % (train_loss / (batch_idx + 1))), - "acc": float("%.3f" % (100. * metrics.accuracy_score(train_true, train_pred))), - "acc_avg": float("%.3f" % (100. * metrics.balanced_accuracy_score(train_true, train_pred))), - "time": time_cost + "acc": float("%.3f" % (100.0 * metrics.accuracy_score(train_true, train_pred))), + "acc_avg": float("%.3f" % (100.0 * metrics.balanced_accuracy_score(train_true, train_pred))), + "time": time_cost, } @@ -233,19 +276,23 @@ def validate(net, testloader, criterion, device): test_pred.append(preds.detach().cpu().numpy()) total += label.size(0) correct += preds.eq(label).sum().item() - progress_bar(batch_idx, len(testloader), 'Loss: %.3f | Acc: %.3f%% (%d/%d)' - % (test_loss / (batch_idx + 1), 100. * correct / total, correct, total)) + progress_bar( + batch_idx, + len(testloader), + "Loss: %.3f | Acc: %.3f%% (%d/%d)" + % (test_loss / (batch_idx + 1), 100.0 * correct / total, correct, total), + ) time_cost = int((datetime.datetime.now() - time_cost).total_seconds()) test_true = np.concatenate(test_true) test_pred = np.concatenate(test_pred) return { "loss": float("%.3f" % (test_loss / (batch_idx + 1))), - "acc": float("%.3f" % (100. * metrics.accuracy_score(test_true, test_pred))), - "acc_avg": float("%.3f" % (100. * metrics.balanced_accuracy_score(test_true, test_pred))), - "time": time_cost + "acc": float("%.3f" % (100.0 * metrics.accuracy_score(test_true, test_pred))), + "acc_avg": float("%.3f" % (100.0 * metrics.balanced_accuracy_score(test_true, test_pred))), + "time": time_cost, } -if __name__ == '__main__': +if __name__ == "__main__": main() diff --git a/classification_ScanObjectNN/models/__init__.py b/classification_ScanObjectNN/models/__init__.py index b8815da..5a546cf 100644 --- a/classification_ScanObjectNN/models/__init__.py +++ b/classification_ScanObjectNN/models/__init__.py @@ -1,3 +1 @@ -from __future__ import absolute_import - from .pointmlp import pointMLP, pointMLPElite diff --git a/classification_ScanObjectNN/models/pointmlp.py b/classification_ScanObjectNN/models/pointmlp.py index 597ba46..8b3e236 100644 --- a/classification_ScanObjectNN/models/pointmlp.py +++ b/classification_ScanObjectNN/models/pointmlp.py @@ -1,35 +1,32 @@ - import torch import torch.nn as nn import torch.nn.functional as F + # from torch import einsum # from einops import rearrange, repeat - - from pointnet2_ops import pointnet2_utils def get_activation(activation): - if activation.lower() == 'gelu': + if activation.lower() == "gelu": return nn.GELU() - elif activation.lower() == 'rrelu': + elif activation.lower() == "rrelu": return nn.RReLU(inplace=True) - elif activation.lower() == 'selu': + elif activation.lower() == "selu": return nn.SELU(inplace=True) - elif activation.lower() == 'silu': + elif activation.lower() == "silu": return nn.SiLU(inplace=True) - elif activation.lower() == 'hardswish': + elif activation.lower() == "hardswish": return nn.Hardswish(inplace=True) - elif activation.lower() == 'leakyrelu': + elif activation.lower() == "leakyrelu": return nn.LeakyReLU(inplace=True) else: return nn.ReLU(inplace=True) def square_distance(src, dst): - """ - Calculate Euclid distance between each two points. - src^T * dst = xn * xm + yn * ym + zn * zm； + """Calculate Euclid distance between each two points. + src^T * dst = xn * xm + yn * ym + zn * zm; sum(src^2, dim=-1) = xn*xn + yn*yn + zn*zn; sum(dst^2, dim=-1) = xm*xm + ym*ym + zm*zm; dist = (xn - xm)^2 + (yn - ym)^2 + (zn - zm)^2 @@ -38,23 +35,23 @@ def square_distance(src, dst): src: source points, [B, N, C] dst: target points, [B, M, C] Output: - dist: per-point square distance, [B, N, M] + dist: per-point square distance, [B, N, M]. """ B, N, _ = src.shape _, M, _ = dst.shape dist = -2 * torch.matmul(src, dst.permute(0, 2, 1)) - dist += torch.sum(src ** 2, -1).view(B, N, 1) - dist += torch.sum(dst ** 2, -1).view(B, 1, M) + dist += torch.sum(src**2, -1).view(B, N, 1) + dist += torch.sum(dst**2, -1).view(B, 1, M) return dist def index_points(points, idx): - """ - Input: + """Input: points: input points data, [B, N, C] - idx: sample index data, [B, S] + idx: sample index data, [B, S]. + Return: - new_points:, indexed points data, [B, S, C] + new_points:, indexed points data, [B, S, C]. """ device = points.device B = points.shape[0] @@ -63,17 +60,15 @@ def index_points(points, idx): repeat_shape = list(idx.shape) repeat_shape[0] = 1 batch_indices = torch.arange(B, dtype=torch.long).to(device).view(view_shape).repeat(repeat_shape) - new_points = points[batch_indices, idx, :] - return new_points + return points[batch_indices, idx, :] def farthest_point_sample(xyz, npoint): - """ - Input: + """Input: xyz: pointcloud data, [B, N, 3] npoint: number of samples Return: - centroids: sampled pointcloud index, [B, npoint] + centroids: sampled pointcloud index, [B, npoint]. """ device = xyz.device B, N, C = xyz.shape @@ -91,21 +86,21 @@ def farthest_point_sample(xyz, npoint): def query_ball_point(radius, nsample, xyz, new_xyz): - """ - Input: + """Input: radius: local region radius nsample: max sample number in local region xyz: all points, [B, N, 3] - new_xyz: query points, [B, S, 3] + new_xyz: query points, [B, S, 3]. + Return: - group_idx: grouped points index, [B, S, nsample] + group_idx: grouped points index, [B, S, nsample]. """ device = xyz.device B, N, C = xyz.shape _, S, _ = new_xyz.shape group_idx = torch.arange(N, dtype=torch.long).to(device).view(1, 1, N).repeat([B, S, 1]) sqrdists = square_distance(new_xyz, xyz) - group_idx[sqrdists > radius ** 2] = N + group_idx[sqrdists > radius**2] = N group_idx = group_idx.sort(dim=-1)[0][:, :, :nsample] group_first = group_idx[:, :, 0].view(B, S, 1).repeat([1, 1, nsample]) mask = group_idx == N @@ -114,13 +109,13 @@ def query_ball_point(radius, nsample, xyz, new_xyz): def knn_point(nsample, xyz, new_xyz): - """ - Input: + """Input: nsample: max sample number in local region xyz: all points, [B, N, C] - new_xyz: query points, [B, S, C] + new_xyz: query points, [B, S, C]. + Return: - group_idx: grouped points index, [B, S, nsample] + group_idx: grouped points index, [B, S, nsample]. """ sqrdists = square_distance(new_xyz, xyz) _, group_idx = torch.topk(sqrdists, nsample, dim=-1, largest=False, sorted=False) @@ -129,13 +124,12 @@ def knn_point(nsample, xyz, new_xyz): class LocalGrouper(nn.Module): def __init__(self, channel, groups, kneighbors, use_xyz=True, normalize="center", **kwargs): - """ - Give xyz[b,p,3] and fea[b,p,d], return new_xyz[b,g,3] and new_fea[b,g,k,d] + """Give xyz[b,p,3] and fea[b,p,d], return new_xyz[b,g,3] and new_fea[b,g,k,d] :param groups: groups number :param kneighbors: k-nerighbors - :param kwargs: others + :param kwargs: others. """ - super(LocalGrouper, self).__init__() + super().__init__() self.groups = groups self.kneighbors = kneighbors self.use_xyz = use_xyz @@ -144,11 +138,11 @@ class LocalGrouper(nn.Module): else: self.normalize = None if self.normalize not in ["center", "anchor"]: - print(f"Unrecognized normalize parameter (self.normalize), set to None. Should be one of [center, anchor].") + print("Unrecognized normalize parameter (self.normalize), set to None. Should be one of [center, anchor].") self.normalize = None if self.normalize is not None: - add_channel=3 if self.use_xyz else 0 - self.affine_alpha = nn.Parameter(torch.ones([1,1,1,channel + add_channel])) + add_channel = 3 if self.use_xyz else 0 + self.affine_alpha = nn.Parameter(torch.ones([1, 1, 1, channel + add_channel])) self.affine_beta = nn.Parameter(torch.zeros([1, 1, 1, channel + add_channel])) def forward(self, xyz, points): @@ -167,29 +161,33 @@ class LocalGrouper(nn.Module): grouped_xyz = index_points(xyz, idx) # [B, npoint, k, 3] grouped_points = index_points(points, idx) # [B, npoint, k, d] if self.use_xyz: - grouped_points = torch.cat([grouped_points, grouped_xyz],dim=-1) # [B, npoint, k, d+3] + grouped_points = torch.cat([grouped_points, grouped_xyz], dim=-1) # [B, npoint, k, d+3] if self.normalize is not None: - if self.normalize =="center": + if self.normalize == "center": mean = torch.mean(grouped_points, dim=2, keepdim=True) - if self.normalize =="anchor": - mean = torch.cat([new_points, new_xyz],dim=-1) if self.use_xyz else new_points + if self.normalize == "anchor": + mean = torch.cat([new_points, new_xyz], dim=-1) if self.use_xyz else new_points mean = mean.unsqueeze(dim=-2) # [B, npoint, 1, d+3] - std = torch.std((grouped_points-mean).reshape(B,-1),dim=-1,keepdim=True).unsqueeze(dim=-1).unsqueeze(dim=-1) - grouped_points = (grouped_points-mean)/(std + 1e-5) - grouped_points = self.affine_alpha*grouped_points + self.affine_beta + std = ( + torch.std((grouped_points - mean).reshape(B, -1), dim=-1, keepdim=True) + .unsqueeze(dim=-1) + .unsqueeze(dim=-1) + ) + grouped_points = (grouped_points - mean) / (std + 1e-5) + grouped_points = self.affine_alpha * grouped_points + self.affine_beta new_points = torch.cat([grouped_points, new_points.view(B, S, 1, -1).repeat(1, 1, self.kneighbors, 1)], dim=-1) return new_xyz, new_points class ConvBNReLU1D(nn.Module): - def __init__(self, in_channels, out_channels, kernel_size=1, bias=True, activation='relu'): - super(ConvBNReLU1D, self).__init__() + def __init__(self, in_channels, out_channels, kernel_size=1, bias=True, activation="relu"): + super().__init__() self.act = get_activation(activation) self.net = nn.Sequential( nn.Conv1d(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, bias=bias), nn.BatchNorm1d(out_channels), - self.act + self.act, ) def forward(self, x): @@ -197,30 +195,43 @@ class ConvBNReLU1D(nn.Module): class ConvBNReLURes1D(nn.Module): - def __init__(self, channel, kernel_size=1, groups=1, res_expansion=1.0, bias=True, activation='relu'): - super(ConvBNReLURes1D, self).__init__() + def __init__(self, channel, kernel_size=1, groups=1, res_expansion=1.0, bias=True, activation="relu"): + super().__init__() self.act = get_activation(activation) self.net1 = nn.Sequential( - nn.Conv1d(in_channels=channel, out_channels=int(channel * res_expansion), - kernel_size=kernel_size, groups=groups, bias=bias), + nn.Conv1d( + in_channels=channel, + out_channels=int(channel * res_expansion), + kernel_size=kernel_size, + groups=groups, + bias=bias, + ), nn.BatchNorm1d(int(channel * res_expansion)), - self.act + self.act, ) if groups > 1: self.net2 = nn.Sequential( - nn.Conv1d(in_channels=int(channel * res_expansion), out_channels=channel, - kernel_size=kernel_size, groups=groups, bias=bias), + nn.Conv1d( + in_channels=int(channel * res_expansion), + out_channels=channel, + kernel_size=kernel_size, + groups=groups, + bias=bias, + ), nn.BatchNorm1d(channel), self.act, - nn.Conv1d(in_channels=channel, out_channels=channel, - kernel_size=kernel_size, bias=bias), + nn.Conv1d(in_channels=channel, out_channels=channel, kernel_size=kernel_size, bias=bias), nn.BatchNorm1d(channel), ) else: self.net2 = nn.Sequential( - nn.Conv1d(in_channels=int(channel * res_expansion), out_channels=channel, - kernel_size=kernel_size, bias=bias), - nn.BatchNorm1d(channel) + nn.Conv1d( + in_channels=int(channel * res_expansion), + out_channels=channel, + kernel_size=kernel_size, + bias=bias, + ), + nn.BatchNorm1d(channel), ) def forward(self, x): @@ -228,21 +239,34 @@ class ConvBNReLURes1D(nn.Module): class PreExtraction(nn.Module): - def __init__(self, channels, out_channels, blocks=1, groups=1, res_expansion=1, bias=True, - activation='relu', use_xyz=True): - """ - input: [b,g,k,d]: output:[b,d,g] + def __init__( + self, + channels, + out_channels, + blocks=1, + groups=1, + res_expansion=1, + bias=True, + activation="relu", + use_xyz=True, + ): + """input: [b,g,k,d]: output:[b,d,g] :param channels: :param blocks: """ - super(PreExtraction, self).__init__() - in_channels = 3+2*channels if use_xyz else 2*channels + super().__init__() + in_channels = 3 + 2 * channels if use_xyz else 2 * channels self.transfer = ConvBNReLU1D(in_channels, out_channels, bias=bias, activation=activation) operation = [] for _ in range(blocks): operation.append( - ConvBNReLURes1D(out_channels, groups=groups, res_expansion=res_expansion, - bias=bias, activation=activation) + ConvBNReLURes1D( + out_channels, + groups=groups, + res_expansion=res_expansion, + bias=bias, + activation=activation, + ), ) self.operation = nn.Sequential(*operation) @@ -254,22 +278,20 @@ class PreExtraction(nn.Module): batch_size, _, _ = x.size() x = self.operation(x) # [b, d, k] x = F.adaptive_max_pool1d(x, 1).view(batch_size, -1) - x = x.reshape(b, n, -1).permute(0, 2, 1) - return x + return x.reshape(b, n, -1).permute(0, 2, 1) class PosExtraction(nn.Module): - def __init__(self, channels, blocks=1, groups=1, res_expansion=1, bias=True, activation='relu'): - """ - input[b,d,g]; output[b,d,g] + def __init__(self, channels, blocks=1, groups=1, res_expansion=1, bias=True, activation="relu"): + """input[b,d,g]; output[b,d,g] :param channels: :param blocks: """ - super(PosExtraction, self).__init__() + super().__init__() operation = [] for _ in range(blocks): operation.append( - ConvBNReLURes1D(channels, groups=groups, res_expansion=res_expansion, bias=bias, activation=activation) + ConvBNReLURes1D(channels, groups=groups, res_expansion=res_expansion, bias=bias, activation=activation), ) self.operation = nn.Sequential(*operation) @@ -278,17 +300,32 @@ class PosExtraction(nn.Module): class Model(nn.Module): - def __init__(self, points=1024, class_num=40, embed_dim=64, groups=1, res_expansion=1.0, - activation="relu", bias=True, use_xyz=True, normalize="center", - dim_expansion=[2, 2, 2, 2], pre_blocks=[2, 2, 2, 2], pos_blocks=[2, 2, 2, 2], - k_neighbors=[32, 32, 32, 32], reducers=[2, 2, 2, 2], **kwargs): - super(Model, self).__init__() + def __init__( + self, + points=1024, + class_num=40, + embed_dim=64, + groups=1, + res_expansion=1.0, + activation="relu", + bias=True, + use_xyz=True, + normalize="center", + dim_expansion=[2, 2, 2, 2], + pre_blocks=[2, 2, 2, 2], + pos_blocks=[2, 2, 2, 2], + k_neighbors=[32, 32, 32, 32], + reducers=[2, 2, 2, 2], + **kwargs, + ): + super().__init__() self.stages = len(pre_blocks) self.class_num = class_num self.points = points self.embedding = ConvBNReLU1D(3, embed_dim, bias=bias, activation=activation) - assert len(pre_blocks) == len(k_neighbors) == len(reducers) == len(pos_blocks) == len(dim_expansion), \ - "Please check stage number consistent for pre_blocks, pos_blocks k_neighbors, reducers." + assert ( + len(pre_blocks) == len(k_neighbors) == len(reducers) == len(pos_blocks) == len(dim_expansion) + ), "Please check stage number consistent for pre_blocks, pos_blocks k_neighbors, reducers." self.local_grouper_list = nn.ModuleList() self.pre_blocks_list = nn.ModuleList() self.pos_blocks_list = nn.ModuleList() @@ -305,13 +342,26 @@ class Model(nn.Module): local_grouper = LocalGrouper(last_channel, anchor_points, kneighbor, use_xyz, normalize) # [b,g,k,d] self.local_grouper_list.append(local_grouper) # append pre_block_list - pre_block_module = PreExtraction(last_channel, out_channel, pre_block_num, groups=groups, - res_expansion=res_expansion, - bias=bias, activation=activation, use_xyz=use_xyz) + pre_block_module = PreExtraction( + last_channel, + out_channel, + pre_block_num, + groups=groups, + res_expansion=res_expansion, + bias=bias, + activation=activation, + use_xyz=use_xyz, + ) self.pre_blocks_list.append(pre_block_module) # append pos_block_list - pos_block_module = PosExtraction(out_channel, pos_block_num, groups=groups, - res_expansion=res_expansion, bias=bias, activation=activation) + pos_block_module = PosExtraction( + out_channel, + pos_block_num, + groups=groups, + res_expansion=res_expansion, + bias=bias, + activation=activation, + ) self.pos_blocks_list.append(pos_block_module) last_channel = out_channel @@ -326,7 +376,7 @@ class Model(nn.Module): nn.BatchNorm1d(256), self.act, nn.Dropout(0.5), - nn.Linear(256, self.class_num) + nn.Linear(256, self.class_num), ) def forward(self, x): @@ -340,29 +390,52 @@ class Model(nn.Module): x = self.pos_blocks_list[i](x) # [b,d,g] x = F.adaptive_max_pool1d(x, 1).squeeze(dim=-1) - x = self.classifier(x) - return x - - + return self.classifier(x) def pointMLP(num_classes=40, **kwargs) -> Model: - return Model(points=1024, class_num=num_classes, embed_dim=64, groups=1, res_expansion=1.0, - activation="relu", bias=False, use_xyz=False, normalize="anchor", - dim_expansion=[2, 2, 2, 2], pre_blocks=[2, 2, 2, 2], pos_blocks=[2, 2, 2, 2], - k_neighbors=[24, 24, 24, 24], reducers=[2, 2, 2, 2], **kwargs) + return Model( + points=1024, + class_num=num_classes, + embed_dim=64, + groups=1, + res_expansion=1.0, + activation="relu", + bias=False, + use_xyz=False, + normalize="anchor", + dim_expansion=[2, 2, 2, 2], + pre_blocks=[2, 2, 2, 2], + pos_blocks=[2, 2, 2, 2], + k_neighbors=[24, 24, 24, 24], + reducers=[2, 2, 2, 2], + **kwargs, + ) def pointMLPElite(num_classes=40, **kwargs) -> Model: - return Model(points=1024, class_num=num_classes, embed_dim=32, groups=1, res_expansion=0.25, - activation="relu", bias=False, use_xyz=False, normalize="anchor", - dim_expansion=[2, 2, 2, 1], pre_blocks=[1, 1, 2, 1], pos_blocks=[1, 1, 2, 1], - k_neighbors=[24,24,24,24], reducers=[2, 2, 2, 2], **kwargs) + return Model( + points=1024, + class_num=num_classes, + embed_dim=32, + groups=1, + res_expansion=0.25, + activation="relu", + bias=False, + use_xyz=False, + normalize="anchor", + dim_expansion=[2, 2, 2, 1], + pre_blocks=[1, 1, 2, 1], + pos_blocks=[1, 1, 2, 1], + k_neighbors=[24, 24, 24, 24], + reducers=[2, 2, 2, 2], + **kwargs, + ) -if __name__ == '__main__': + +if __name__ == "__main__": data = torch.rand(2, 3, 1024) print("===> testing pointMLP ...") model = pointMLP() out = model(data) print(out.shape) - diff --git a/classification_ScanObjectNN/utils/__init__.py b/classification_ScanObjectNN/utils/__init__.py index 09e4b76..9be60ee 100644 --- a/classification_ScanObjectNN/utils/__init__.py +++ b/classification_ScanObjectNN/utils/__init__.py @@ -1,5 +1,4 @@ -"""Useful utils -""" -from .misc import * +"""Useful utils.""" from .logger import * +from .misc import * from .progress.progress.bar import Bar as Bar diff --git a/classification_ScanObjectNN/utils/logger.py b/classification_ScanObjectNN/utils/logger.py index 7eb5c67..db15f5f 100644 --- a/classification_ScanObjectNN/utils/logger.py +++ b/classification_ScanObjectNN/utils/logger.py @@ -1,89 +1,92 @@ # A simple torch style logger # (C) Wei YANG 2017 -from __future__ import absolute_import + import matplotlib.pyplot as plt -import os -import sys import numpy as np -__all__ = ['Logger', 'LoggerMonitor', 'savefig'] +__all__ = ["Logger", "LoggerMonitor", "savefig"] + def savefig(fname, dpi=None): - dpi = 150 if dpi == None else dpi + dpi = 150 if dpi is None else dpi plt.savefig(fname, dpi=dpi) - + + def plot_overlap(logger, names=None): - names = logger.names if names == None else names + names = logger.names if names is None else names numbers = logger.numbers for _, name in enumerate(names): x = np.arange(len(numbers[name])) plt.plot(x, np.asarray(numbers[name])) - return [logger.title + '(' + name + ')' for name in names] + return [logger.title + "(" + name + ")" for name in names] -class Logger(object): - '''Save training process to log file with simple plot function.''' - def __init__(self, fpath, title=None, resume=False): + +class Logger: + """Save training process to log file with simple plot function.""" + + def __init__(self, fpath, title=None, resume=False): self.file = None self.resume = resume - self.title = '' if title == None else title + self.title = "" if title is None else title if fpath is not None: - if resume: - self.file = open(fpath, 'r') + if resume: + self.file = open(fpath) name = self.file.readline() - self.names = name.rstrip().split('\t') + self.names = name.rstrip().split("\t") self.numbers = {} for _, name in enumerate(self.names): self.numbers[name] = [] for numbers in self.file: - numbers = numbers.rstrip().split('\t') + numbers = numbers.rstrip().split("\t") for i in range(0, len(numbers)): self.numbers[self.names[i]].append(numbers[i]) self.file.close() - self.file = open(fpath, 'a') + self.file = open(fpath, "a") else: - self.file = open(fpath, 'w') + self.file = open(fpath, "w") def set_names(self, names): - if self.resume: + if self.resume: pass # initialize numbers as empty list self.numbers = {} self.names = names for _, name in enumerate(self.names): self.file.write(name) - self.file.write('\t') + self.file.write("\t") self.numbers[name] = [] - self.file.write('\n') + self.file.write("\n") self.file.flush() - def append(self, numbers): - assert len(self.names) == len(numbers), 'Numbers do not match names' + assert len(self.names) == len(numbers), "Numbers do not match names" for index, num in enumerate(numbers): - self.file.write("{0:.6f}".format(num)) - self.file.write('\t') + self.file.write(f"{num:.6f}") + self.file.write("\t") self.numbers[self.names[index]].append(num) - self.file.write('\n') + self.file.write("\n") self.file.flush() - def plot(self, names=None): - names = self.names if names == None else names + def plot(self, names=None): + names = self.names if names is None else names numbers = self.numbers for _, name in enumerate(names): x = np.arange(len(numbers[name])) plt.plot(x, np.asarray(numbers[name])) - plt.legend([self.title + '(' + name + ')' for name in names]) + plt.legend([self.title + "(" + name + ")" for name in names]) plt.grid(True) def close(self): if self.file is not None: self.file.close() -class LoggerMonitor(object): - '''Load and visualize multiple logs.''' - def __init__ (self, paths): - '''paths is a distionary with {name:filepath} pair''' + +class LoggerMonitor: + """Load and visualize multiple logs.""" + + def __init__(self, paths): + """Paths is a distionary with {name:filepath} pair.""" self.loggers = [] for title, path in paths.items(): logger = Logger(path, title=title, resume=True) @@ -95,10 +98,11 @@ class LoggerMonitor(object): legend_text = [] for logger in self.loggers: legend_text += plot_overlap(logger, names) - plt.legend(legend_text, bbox_to_anchor=(1.05, 1), loc=2, borderaxespad=0.) + plt.legend(legend_text, bbox_to_anchor=(1.05, 1), loc=2, borderaxespad=0.0) plt.grid(True) - -if __name__ == '__main__': + + +if __name__ == "__main__": # # Example # logger = Logger('test.txt') # logger.set_names(['Train loss', 'Valid loss','Test loss']) @@ -115,13 +119,13 @@ if __name__ == '__main__': # Example: logger monitor paths = { - 'resadvnet20':'/home/wyang/code/pytorch-classification/checkpoint/cifar10/resadvnet20/log.txt', - 'resadvnet32':'/home/wyang/code/pytorch-classification/checkpoint/cifar10/resadvnet32/log.txt', - 'resadvnet44':'/home/wyang/code/pytorch-classification/checkpoint/cifar10/resadvnet44/log.txt', + "resadvnet20": "/home/wyang/code/pytorch-classification/checkpoint/cifar10/resadvnet20/log.txt", + "resadvnet32": "/home/wyang/code/pytorch-classification/checkpoint/cifar10/resadvnet32/log.txt", + "resadvnet44": "/home/wyang/code/pytorch-classification/checkpoint/cifar10/resadvnet44/log.txt", } - field = ['Valid Acc.'] + field = ["Valid Acc."] monitor = LoggerMonitor(paths) monitor.plot(names=field) - savefig('test.eps') \ No newline at end of file + savefig("test.eps") diff --git a/classification_ScanObjectNN/utils/misc.py b/classification_ScanObjectNN/utils/misc.py index 35f743a..3d073b8 100644 --- a/classification_ScanObjectNN/utils/misc.py +++ b/classification_ScanObjectNN/utils/misc.py @@ -1,48 +1,56 @@ -'''Some helper functions for PyTorch, including: - - get_mean_and_std: calculate the mean and std value of dataset. - - msr_init: net parameter initialization. - - progress_bar: progress bar mimic xlua.progress. -''' +"""Some helper functions for PyTorch, including: +- get_mean_and_std: calculate the mean and std value of dataset. +- msr_init: net parameter initialization. +- progress_bar: progress bar mimic xlua.progress. +""" import errno import os +import random +import shutil import sys import time -import math -import torch -import shutil + import numpy as np -import random -import torch.nn.functional as F - - +import torch import torch.nn as nn +import torch.nn.functional as F import torch.nn.init as init -from torch.autograd import Variable -__all__ = ['get_mean_and_std', 'init_params', 'mkdir_p', 'AverageMeter', - 'progress_bar','save_model',"save_args","set_seed", "IOStream", "cal_loss"] +__all__ = [ + "get_mean_and_std", + "init_params", + "mkdir_p", + "AverageMeter", + "progress_bar", + "save_model", + "save_args", + "set_seed", + "IOStream", + "cal_loss", +] def get_mean_and_std(dataset): - '''Compute the mean and std value of dataset.''' - dataloader = trainloader = torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=True, num_workers=2) + """Compute the mean and std value of dataset.""" + dataloader = torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=True, num_workers=2) mean = torch.zeros(3) std = torch.zeros(3) - print('==> Computing mean and std..') - for inputs, targets in dataloader: + print("==> Computing mean and std..") + for inputs, _targets in dataloader: for i in range(3): - mean[i] += inputs[:,i,:,:].mean() - std[i] += inputs[:,i,:,:].std() + mean[i] += inputs[:, i, :, :].mean() + std[i] += inputs[:, i, :, :].std() mean.div_(len(dataset)) std.div_(len(dataset)) return mean, std + def init_params(net): - '''Init layer parameters.''' + """Init layer parameters.""" for m in net.modules(): if isinstance(m, nn.Conv2d): - init.kaiming_normal(m.weight, mode='fan_out') + init.kaiming_normal(m.weight, mode="fan_out") if m.bias: init.constant(m.bias, 0) elif isinstance(m, nn.BatchNorm2d): @@ -53,8 +61,9 @@ def init_params(net): if m.bias: init.constant(m.bias, 0) + def mkdir_p(path): - '''make dir if not exist''' + """Make dir if not exist.""" try: os.makedirs(path) except OSError as exc: # Python >2.5 @@ -63,10 +72,12 @@ def mkdir_p(path): else: raise -class AverageMeter(object): + +class AverageMeter: """Computes and stores the average and current value - Imported from https://github.com/pytorch/examples/blob/master/imagenet/main.py#L247-L262 + Imported from https://github.com/pytorch/examples/blob/master/imagenet/main.py#L247-L262. """ + def __init__(self): self.reset() @@ -83,25 +94,26 @@ class AverageMeter(object): self.avg = self.sum / self.count - -TOTAL_BAR_LENGTH = 65. +TOTAL_BAR_LENGTH = 65.0 last_time = time.time() begin_time = last_time + + def progress_bar(current, total, msg=None): global last_time, begin_time if current == 0: begin_time = time.time() # Reset for new bar. - cur_len = int(TOTAL_BAR_LENGTH*current/total) + cur_len = int(TOTAL_BAR_LENGTH * current / total) rest_len = int(TOTAL_BAR_LENGTH - cur_len) - 1 - sys.stdout.write(' [') - for i in range(cur_len): - sys.stdout.write('=') - sys.stdout.write('>') - for i in range(rest_len): - sys.stdout.write('.') - sys.stdout.write(']') + sys.stdout.write(" [") + for _i in range(cur_len): + sys.stdout.write("=") + sys.stdout.write(">") + for _i in range(rest_len): + sys.stdout.write(".") + sys.stdout.write("]") cur_time = time.time() step_time = cur_time - last_time @@ -109,12 +121,12 @@ def progress_bar(current, total, msg=None): tot_time = cur_time - begin_time L = [] - L.append(' Step: %s' % format_time(step_time)) - L.append(' | Tot: %s' % format_time(tot_time)) + L.append(" Step: %s" % format_time(step_time)) + L.append(" | Tot: %s" % format_time(tot_time)) if msg: - L.append(' | ' + msg) + L.append(" | " + msg) - msg = ''.join(L) + msg = "".join(L) sys.stdout.write(msg) # for i in range(term_width-int(TOTAL_BAR_LENGTH)-len(msg)-3): # sys.stdout.write(' ') @@ -122,76 +134,74 @@ def progress_bar(current, total, msg=None): # Go back to the center of the bar. # for i in range(term_width-int(TOTAL_BAR_LENGTH/2)+2): # sys.stdout.write('\b') - sys.stdout.write(' %d/%d ' % (current+1, total)) + sys.stdout.write(" %d/%d " % (current + 1, total)) - if current < total-1: - sys.stdout.write('\r') + if current < total - 1: + sys.stdout.write("\r") else: - sys.stdout.write('\n') + sys.stdout.write("\n") sys.stdout.flush() def format_time(seconds): - days = int(seconds / 3600/24) - seconds = seconds - days*3600*24 + days = int(seconds / 3600 / 24) + seconds = seconds - days * 3600 * 24 hours = int(seconds / 3600) - seconds = seconds - hours*3600 + seconds = seconds - hours * 3600 minutes = int(seconds / 60) - seconds = seconds - minutes*60 + seconds = seconds - minutes * 60 secondsf = int(seconds) seconds = seconds - secondsf - millis = int(seconds*1000) + millis = int(seconds * 1000) - f = '' + f = "" i = 1 if days > 0: - f += str(days) + 'D' + f += str(days) + "D" i += 1 if hours > 0 and i <= 2: - f += str(hours) + 'h' + f += str(hours) + "h" i += 1 if minutes > 0 and i <= 2: - f += str(minutes) + 'm' + f += str(minutes) + "m" i += 1 if secondsf > 0 and i <= 2: - f += str(secondsf) + 's' + f += str(secondsf) + "s" i += 1 if millis > 0 and i <= 2: - f += str(millis) + 'ms' + f += str(millis) + "ms" i += 1 - if f == '': - f = '0ms' + if f == "": + f = "0ms" return f def save_model(net, epoch, path, acc, is_best, **kwargs): state = { - 'net': net.state_dict(), - 'epoch': epoch, - 'acc': acc + "net": net.state_dict(), + "epoch": epoch, + "acc": acc, } for key, value in kwargs.items(): state[key] = value filepath = os.path.join(path, "last_checkpoint.pth") torch.save(state, filepath) if is_best: - shutil.copyfile(filepath, os.path.join(path, 'best_checkpoint.pth')) - + shutil.copyfile(filepath, os.path.join(path, "best_checkpoint.pth")) def save_args(args): - file = open(os.path.join(args.checkpoint, 'args.txt'), "w") + file = open(os.path.join(args.checkpoint, "args.txt"), "w") for k, v in vars(args).items(): file.write(f"{k}:\t {v}\n") file.close() - def set_seed(seed=None): if seed is None: return random.seed(seed) - os.environ['PYTHONHASHSEED'] = ("%s" % seed) + os.environ["PYTHONHASHSEED"] = "%s" % seed np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed(seed) @@ -200,15 +210,14 @@ def set_seed(seed=None): torch.backends.cudnn.deterministic = True - # create a file and write the text into it -class IOStream(): +class IOStream: def __init__(self, path): - self.f = open(path, 'a') + self.f = open(path, "a") def cprint(self, text): print(text) - self.f.write(text+'\n') + self.f.write(text + "\n") self.f.flush() def close(self): @@ -216,8 +225,7 @@ class IOStream(): def cal_loss(pred, gold, smoothing=True): - ''' Calculate cross entropy loss, apply label smoothing if needed. ''' - + """Calculate cross entropy loss, apply label smoothing if needed.""" gold = gold.contiguous().view(-1) if smoothing: @@ -230,6 +238,6 @@ def cal_loss(pred, gold, smoothing=True): loss = -(one_hot * log_prb).sum(dim=1).mean() else: - loss = F.cross_entropy(pred, gold, reduction='mean') + loss = F.cross_entropy(pred, gold, reduction="mean") return loss diff --git a/classification_ScanObjectNN/utils/progress/progress/__init__.py b/classification_ScanObjectNN/utils/progress/progress/__init__.py index 09dfc1e..1166f4b 100644 --- a/classification_ScanObjectNN/utils/progress/progress/__init__.py +++ b/classification_ScanObjectNN/utils/progress/progress/__init__.py @@ -12,7 +12,6 @@ # ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF # OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE. -from __future__ import division from collections import deque from datetime import timedelta @@ -20,13 +19,12 @@ from math import ceil from sys import stderr from time import time - -__version__ = '1.3' +__version__ = "1.3" -class Infinite(object): +class Infinite: file = stderr - sma_window = 10 # Simple Moving Average window + sma_window = 10 # Simple Moving Average window def __init__(self, *args, **kwargs): self.index = 0 @@ -38,7 +36,7 @@ class Infinite(object): setattr(self, key, val) def __getitem__(self, key): - if key.startswith('_'): + if key.startswith("_"): return None return getattr(self, key, None) @@ -83,8 +81,8 @@ class Infinite(object): class Progress(Infinite): def __init__(self, *args, **kwargs): - super(Progress, self).__init__(*args, **kwargs) - self.max = kwargs.get('max', 100) + super().__init__(*args, **kwargs) + self.max = kwargs.get("max", 100) @property def eta(self): diff --git a/classification_ScanObjectNN/utils/progress/progress/bar.py b/classification_ScanObjectNN/utils/progress/progress/bar.py index 5ee968f..93ec858 100644 --- a/classification_ScanObjectNN/utils/progress/progress/bar.py +++ b/classification_ScanObjectNN/utils/progress/progress/bar.py @@ -1,5 +1,3 @@ -# -*- coding: utf-8 -*- - # Copyright (c) 2012 Giorgos Verigakis # # Permission to use, copy, modify, and distribute this software for any @@ -14,19 +12,18 @@ # ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF # OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE. -from __future__ import unicode_literals from . import Progress from .helpers import WritelnMixin class Bar(WritelnMixin, Progress): width = 32 - message = '' - suffix = '%(index)d/%(max)d' - bar_prefix = ' |' - bar_suffix = '| ' - empty_fill = ' ' - fill = '#' + message = "" + suffix = "%(index)d/%(max)d" + bar_prefix = " |" + bar_suffix = "| " + empty_fill = " " + fill = "#" hide_cursor = True def update(self): @@ -37,52 +34,50 @@ class Bar(WritelnMixin, Progress): bar = self.fill * filled_length empty = self.empty_fill * empty_length suffix = self.suffix % self - line = ''.join([message, self.bar_prefix, bar, empty, self.bar_suffix, - suffix]) + line = "".join([message, self.bar_prefix, bar, empty, self.bar_suffix, suffix]) self.writeln(line) class ChargingBar(Bar): - suffix = '%(percent)d%%' - bar_prefix = ' ' - bar_suffix = ' ' - empty_fill = '∙' - fill = '█' + suffix = "%(percent)d%%" + bar_prefix = " " + bar_suffix = " " + empty_fill = "∙" + fill = "█" class FillingSquaresBar(ChargingBar): - empty_fill = '▢' - fill = '▣' + empty_fill = "▢" + fill = "▣" class FillingCirclesBar(ChargingBar): - empty_fill = '◯' - fill = '◉' + empty_fill = "◯" + fill = "◉" class IncrementalBar(Bar): - phases = (' ', '▏', '▎', '▍', '▌', '▋', '▊', '▉', '█') + phases = (" ", "▏", "▎", "▍", "▌", "▋", "▊", "▉", "█") def update(self): nphases = len(self.phases) filled_len = self.width * self.progress - nfull = int(filled_len) # Number of full chars + nfull = int(filled_len) # Number of full chars phase = int((filled_len - nfull) * nphases) # Phase of last char - nempty = self.width - nfull # Number of empty chars + nempty = self.width - nfull # Number of empty chars message = self.message % self bar = self.phases[-1] * nfull - current = self.phases[phase] if phase > 0 else '' + current = self.phases[phase] if phase > 0 else "" empty = self.empty_fill * max(0, nempty - len(current)) suffix = self.suffix % self - line = ''.join([message, self.bar_prefix, bar, current, empty, - self.bar_suffix, suffix]) + line = "".join([message, self.bar_prefix, bar, current, empty, self.bar_suffix, suffix]) self.writeln(line) class PixelBar(IncrementalBar): - phases = ('⡀', '⡄', '⡆', '⡇', '⣇', '⣧', '⣷', '⣿') + phases = ("⡀", "⡄", "⡆", "⡇", "⣇", "⣧", "⣷", "⣿") class ShadyBar(IncrementalBar): - phases = (' ', '░', '▒', '▓', '█') + phases = (" ", "░", "▒", "▓", "█") diff --git a/classification_ScanObjectNN/utils/progress/progress/counter.py b/classification_ScanObjectNN/utils/progress/progress/counter.py index 6b45a1e..90da4ff 100644 --- a/classification_ScanObjectNN/utils/progress/progress/counter.py +++ b/classification_ScanObjectNN/utils/progress/progress/counter.py @@ -1,5 +1,3 @@ -# -*- coding: utf-8 -*- - # Copyright (c) 2012 Giorgos Verigakis # # Permission to use, copy, modify, and distribute this software for any @@ -14,13 +12,12 @@ # ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF # OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE. -from __future__ import unicode_literals from . import Infinite, Progress from .helpers import WriteMixin class Counter(WriteMixin, Infinite): - message = '' + message = "" hide_cursor = True def update(self): @@ -35,7 +32,7 @@ class Countdown(WriteMixin, Progress): class Stack(WriteMixin, Progress): - phases = (' ', '▁', '▂', '▃', '▄', '▅', '▆', '▇', '█') + phases = (" ", "▁", "▂", "▃", "▄", "▅", "▆", "▇", "█") hide_cursor = True def update(self): @@ -45,4 +42,4 @@ class Stack(WriteMixin, Progress): class Pie(Stack): - phases = ('○', '◔', '◑', '◕', '●') + phases = ("○", "◔", "◑", "◕", "●") diff --git a/classification_ScanObjectNN/utils/progress/progress/helpers.py b/classification_ScanObjectNN/utils/progress/progress/helpers.py index 9ed90b2..04abe24 100644 --- a/classification_ScanObjectNN/utils/progress/progress/helpers.py +++ b/classification_ScanObjectNN/utils/progress/progress/helpers.py @@ -12,78 +12,76 @@ # ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF # OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE. -from __future__ import print_function + +HIDE_CURSOR = "\x1b[?25l" +SHOW_CURSOR = "\x1b[?25h" -HIDE_CURSOR = '\x1b[?25l' -SHOW_CURSOR = '\x1b[?25h' - - -class WriteMixin(object): +class WriteMixin: hide_cursor = False def __init__(self, message=None, **kwargs): - super(WriteMixin, self).__init__(**kwargs) + super().__init__(**kwargs) self._width = 0 if message: self.message = message if self.file.isatty(): if self.hide_cursor: - print(HIDE_CURSOR, end='', file=self.file) - print(self.message, end='', file=self.file) + print(HIDE_CURSOR, end="", file=self.file) + print(self.message, end="", file=self.file) self.file.flush() def write(self, s): if self.file.isatty(): - b = '\b' * self._width + b = "\b" * self._width c = s.ljust(self._width) - print(b + c, end='', file=self.file) + print(b + c, end="", file=self.file) self._width = max(self._width, len(s)) self.file.flush() def finish(self): if self.file.isatty() and self.hide_cursor: - print(SHOW_CURSOR, end='', file=self.file) + print(SHOW_CURSOR, end="", file=self.file) -class WritelnMixin(object): +class WritelnMixin: hide_cursor = False def __init__(self, message=None, **kwargs): - super(WritelnMixin, self).__init__(**kwargs) + super().__init__(**kwargs) if message: self.message = message if self.file.isatty() and self.hide_cursor: - print(HIDE_CURSOR, end='', file=self.file) + print(HIDE_CURSOR, end="", file=self.file) def clearln(self): if self.file.isatty(): - print('\r\x1b[K', end='', file=self.file) + print("\r\x1b[K", end="", file=self.file) def writeln(self, line): if self.file.isatty(): self.clearln() - print(line, end='', file=self.file) + print(line, end="", file=self.file) self.file.flush() def finish(self): if self.file.isatty(): print(file=self.file) if self.hide_cursor: - print(SHOW_CURSOR, end='', file=self.file) + print(SHOW_CURSOR, end="", file=self.file) -from signal import signal, SIGINT +from signal import SIGINT, signal from sys import exit -class SigIntMixin(object): - """Registers a signal handler that calls finish on SIGINT""" +class SigIntMixin: + """Registers a signal handler that calls finish on SIGINT.""" def __init__(self, *args, **kwargs): - super(SigIntMixin, self).__init__(*args, **kwargs) + super().__init__(*args, **kwargs) signal(SIGINT, self._sigint_handler) def _sigint_handler(self, signum, frame): diff --git a/classification_ScanObjectNN/utils/progress/progress/spinner.py b/classification_ScanObjectNN/utils/progress/progress/spinner.py index 464c7b2..77d0f49 100644 --- a/classification_ScanObjectNN/utils/progress/progress/spinner.py +++ b/classification_ScanObjectNN/utils/progress/progress/spinner.py @@ -1,5 +1,3 @@ -# -*- coding: utf-8 -*- - # Copyright (c) 2012 Giorgos Verigakis # # Permission to use, copy, modify, and distribute this software for any @@ -14,14 +12,13 @@ # ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF # OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE. -from __future__ import unicode_literals from . import Infinite from .helpers import WriteMixin class Spinner(WriteMixin, Infinite): - message = '' - phases = ('-', '\\', '|', '/') + message = "" + phases = ("-", "\\", "|", "/") hide_cursor = True def update(self): @@ -30,15 +27,16 @@ class Spinner(WriteMixin, Infinite): class PieSpinner(Spinner): - phases = ['◷', '◶', '◵', '◴'] + phases = ["◷", "◶", "◵", "◴"] class MoonSpinner(Spinner): - phases = ['◑', '◒', '◐', '◓'] + phases = ["◑", "◒", "◐", "◓"] class LineSpinner(Spinner): - phases = ['⎺', '⎻', '⎼', '⎽', '⎼', '⎻'] + phases = ["⎺", "⎻", "⎼", "⎽", "⎼", "⎻"] + class PixelSpinner(Spinner): - phases = ['⣾','⣷', '⣯', '⣟', '⡿', '⢿', '⣻', '⣽'] + phases = ["⣾", "⣷", "⣯", "⣟", "⡿", "⢿", "⣻", "⣽"] diff --git a/classification_ScanObjectNN/utils/progress/setup.py b/classification_ScanObjectNN/utils/progress/setup.py index c877781..4d935fd 100755 --- a/classification_ScanObjectNN/utils/progress/setup.py +++ b/classification_ScanObjectNN/utils/progress/setup.py @@ -1,29 +1,27 @@ #!/usr/bin/env python +import progress from setuptools import setup -import progress - - setup( - name='progress', + name="progress", version=progress.__version__, - description='Easy to use progress bars', - long_description=open('README.rst').read(), - author='Giorgos Verigakis', - author_email='verigak@gmail.com', - url='http://github.com/verigak/progress/', - license='ISC', - packages=['progress'], + description="Easy to use progress bars", + long_description=open("README.rst").read(), + author="Giorgos Verigakis", + author_email="verigak@gmail.com", + url="http://github.com/verigak/progress/", + license="ISC", + packages=["progress"], classifiers=[ - 'Environment :: Console', - 'Intended Audience :: Developers', - 'License :: OSI Approved :: ISC License (ISCL)', - 'Programming Language :: Python :: 2.6', - 'Programming Language :: Python :: 2.7', - 'Programming Language :: Python :: 3.3', - 'Programming Language :: Python :: 3.4', - 'Programming Language :: Python :: 3.5', - 'Programming Language :: Python :: 3.6', - ] + "Environment :: Console", + "Intended Audience :: Developers", + "License :: OSI Approved :: ISC License (ISCL)", + "Programming Language :: Python :: 2.6", + "Programming Language :: Python :: 2.7", + "Programming Language :: Python :: 3.3", + "Programming Language :: Python :: 3.4", + "Programming Language :: Python :: 3.5", + "Programming Language :: Python :: 3.6", + ], ) diff --git a/classification_ScanObjectNN/utils/progress/test_progress.py b/classification_ScanObjectNN/utils/progress/test_progress.py index 0f68b01..1f200c3 100755 --- a/classification_ScanObjectNN/utils/progress/test_progress.py +++ b/classification_ScanObjectNN/utils/progress/test_progress.py @@ -1,16 +1,12 @@ #!/usr/bin/env python -from __future__ import print_function import random import time -from progress.bar import (Bar, ChargingBar, FillingSquaresBar, - FillingCirclesBar, IncrementalBar, PixelBar, - ShadyBar) -from progress.spinner import (Spinner, PieSpinner, MoonSpinner, LineSpinner, - PixelSpinner) -from progress.counter import Counter, Countdown, Stack, Pie +from progress.bar import Bar, ChargingBar, FillingCirclesBar, FillingSquaresBar, IncrementalBar, PixelBar, ShadyBar +from progress.counter import Countdown, Counter, Pie, Stack +from progress.spinner import LineSpinner, MoonSpinner, PieSpinner, PixelSpinner, Spinner def sleep(): @@ -20,29 +16,29 @@ def sleep(): for bar_cls in (Bar, ChargingBar, FillingSquaresBar, FillingCirclesBar): - suffix = '%(index)d/%(max)d [%(elapsed)d / %(eta)d / %(eta_td)s]' + suffix = "%(index)d/%(max)d [%(elapsed)d / %(eta)d / %(eta_td)s]" bar = bar_cls(bar_cls.__name__, suffix=suffix) - for i in bar.iter(range(200)): + for _i in bar.iter(range(200)): sleep() for bar_cls in (IncrementalBar, PixelBar, ShadyBar): - suffix = '%(percent)d%% [%(elapsed_td)s / %(eta)d / %(eta_td)s]' + suffix = "%(percent)d%% [%(elapsed_td)s / %(eta)d / %(eta_td)s]" bar = bar_cls(bar_cls.__name__, suffix=suffix) - for i in bar.iter(range(200)): + for _i in bar.iter(range(200)): sleep() for spin in (Spinner, PieSpinner, MoonSpinner, LineSpinner, PixelSpinner): - for i in spin(spin.__name__ + ' ').iter(range(100)): + for _i in spin(spin.__name__ + " ").iter(range(100)): sleep() print() for singleton in (Counter, Countdown, Stack, Pie): - for i in singleton(singleton.__name__ + ' ').iter(range(100)): + for _i in singleton(singleton.__name__ + " ").iter(range(100)): sleep() print() -bar = IncrementalBar('Random', suffix='%(index)d') -for i in range(100): +bar = IncrementalBar("Random", suffix="%(index)d") +for _i in range(100): bar.goto(random.randint(0, 100)) sleep() bar.finish() diff --git a/part_segmentation/main.py b/part_segmentation/main.py index d504a45..a08c8fa 100644 --- a/part_segmentation/main.py +++ b/part_segmentation/main.py @@ -1,30 +1,46 @@ -from __future__ import print_function -import os import argparse -import torch -import torch.optim as optim -from torch.optim.lr_scheduler import CosineAnnealingLR, StepLR -from util.data_util import PartNormalDataset -import torch.nn.functional as F -import torch.nn as nn +import os +import random +from collections import defaultdict + import model as models import numpy as np -from torch.utils.data import DataLoader -from util.util import to_categorical, compute_overall_iou, IOStream -from tqdm import tqdm -from collections import defaultdict +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.optim as optim from torch.autograd import Variable -import random +from torch.optim.lr_scheduler import CosineAnnealingLR, StepLR +from torch.utils.data import DataLoader +from tqdm import tqdm +from util.data_util import PartNormalDataset +from util.util import IOStream, compute_overall_iou, to_categorical - -classes_str = ['aero','bag','cap','car','chair','ear','guitar','knife','lamp','lapt','moto','mug','Pistol','rock','stake','table'] +classes_str = [ + "aero", + "bag", + "cap", + "car", + "chair", + "ear", + "guitar", + "knife", + "lamp", + "lapt", + "moto", + "mug", + "Pistol", + "rock", + "stake", + "table", +] def _init_(): - if not os.path.exists('checkpoints'): - os.makedirs('checkpoints') - if not os.path.exists('checkpoints/' + args.exp_name): - os.makedirs('checkpoints/' + args.exp_name) + if not os.path.exists("checkpoints"): + os.makedirs("checkpoints") + if not os.path.exists("checkpoints/" + args.exp_name): + os.makedirs("checkpoints/" + args.exp_name) def weight_init(m): @@ -49,7 +65,6 @@ def weight_init(m): def train(args, io): - # ============= Model =================== num_part = 50 device = torch.device("cuda" if args.cuda else "cpu") @@ -61,16 +76,19 @@ def train(args, io): model = nn.DataParallel(model) print("Let's use", torch.cuda.device_count(), "GPUs!") - '''Resume or not''' + """Resume or not""" if args.resume: - state_dict = torch.load("checkpoints/%s/best_insiou_model.pth" % args.exp_name, - map_location=torch.device('cpu'))['model'] + state_dict = torch.load( + "checkpoints/%s/best_insiou_model.pth" % args.exp_name, + map_location=torch.device("cpu"), + )["model"] for k in state_dict.keys(): - if 'module' not in k: + if "module" not in k: from collections import OrderedDict + new_state_dict = OrderedDict() for k in state_dict: - new_state_dict['module.' + k] = state_dict[k] + new_state_dict["module." + k] = state_dict[k] state_dict = new_state_dict break model.load_state_dict(state_dict) @@ -81,27 +99,37 @@ def train(args, io): print("Training from scratch...") # =========== Dataloader ================= - train_data = PartNormalDataset(npoints=2048, split='trainval', normalize=False) + train_data = PartNormalDataset(npoints=2048, split="trainval", normalize=False) print("The number of training data is:%d", len(train_data)) - test_data = PartNormalDataset(npoints=2048, split='test', normalize=False) + test_data = PartNormalDataset(npoints=2048, split="test", normalize=False) print("The number of test data is:%d", len(test_data)) - train_loader = DataLoader(train_data, batch_size=args.batch_size, shuffle=True, num_workers=args.workers, - drop_last=True) + train_loader = DataLoader( + train_data, + batch_size=args.batch_size, + shuffle=True, + num_workers=args.workers, + drop_last=True, + ) - test_loader = DataLoader(test_data, batch_size=args.test_batch_size, shuffle=False, num_workers=args.workers, - drop_last=False) + test_loader = DataLoader( + test_data, + batch_size=args.test_batch_size, + shuffle=False, + num_workers=args.workers, + drop_last=False, + ) # ============= Optimizer ================ if args.use_sgd: print("Use SGD") - opt = optim.SGD(model.parameters(), lr=args.lr*100, momentum=args.momentum, weight_decay=0) + opt = optim.SGD(model.parameters(), lr=args.lr * 100, momentum=args.momentum, weight_decay=0) else: print("Use Adam") opt = optim.Adam(model.parameters(), lr=args.lr, betas=(0.9, 0.999), eps=1e-08, weight_decay=0) - if args.scheduler == 'cos': + if args.scheduler == "cos": print("Use CosLR") scheduler = CosineAnnealingLR(opt, args.epochs, eta_min=args.lr if args.use_sgd else args.lr / 100) else: @@ -116,28 +144,33 @@ def train(args, io): num_classes = 16 for epoch in range(args.epochs): - train_epoch(train_loader, model, opt, scheduler, epoch, num_part, num_classes, io) test_metrics, total_per_cat_iou = test_epoch(test_loader, model, epoch, num_part, num_classes, io) # 1. when get the best accuracy, save the model: - if test_metrics['accuracy'] > best_acc: - best_acc = test_metrics['accuracy'] - io.cprint('Max Acc:%.5f' % best_acc) + if test_metrics["accuracy"] > best_acc: + best_acc = test_metrics["accuracy"] + io.cprint("Max Acc:%.5f" % best_acc) state = { - 'model': model.module.state_dict() if torch.cuda.device_count() > 1 else model.state_dict(), - 'optimizer': opt.state_dict(), 'epoch': epoch, 'test_acc': best_acc} - torch.save(state, 'checkpoints/%s/best_acc_model.pth' % args.exp_name) + "model": model.module.state_dict() if torch.cuda.device_count() > 1 else model.state_dict(), + "optimizer": opt.state_dict(), + "epoch": epoch, + "test_acc": best_acc, + } + torch.save(state, "checkpoints/%s/best_acc_model.pth" % args.exp_name) # 2. when get the best instance_iou, save the model: - if test_metrics['shape_avg_iou'] > best_instance_iou: - best_instance_iou = test_metrics['shape_avg_iou'] - io.cprint('Max instance iou:%.5f' % best_instance_iou) + if test_metrics["shape_avg_iou"] > best_instance_iou: + best_instance_iou = test_metrics["shape_avg_iou"] + io.cprint("Max instance iou:%.5f" % best_instance_iou) state = { - 'model': model.module.state_dict() if torch.cuda.device_count() > 1 else model.state_dict(), - 'optimizer': opt.state_dict(), 'epoch': epoch, 'test_instance_iou': best_instance_iou} - torch.save(state, 'checkpoints/%s/best_insiou_model.pth' % args.exp_name) + "model": model.module.state_dict() if torch.cuda.device_count() > 1 else model.state_dict(), + "optimizer": opt.state_dict(), + "epoch": epoch, + "test_instance_iou": best_instance_iou, + } + torch.save(state, "checkpoints/%s/best_insiou_model.pth" % args.exp_name) # 3. when get the best class_iou, save the model: # first we need to calculate the average per-class iou @@ -149,22 +182,28 @@ def train(args, io): best_class_iou = avg_class_iou # print the iou of each class: for cat_idx in range(16): - io.cprint(classes_str[cat_idx] + ' iou: ' + str(total_per_cat_iou[cat_idx])) - io.cprint('Max class iou:%.5f' % best_class_iou) + io.cprint(classes_str[cat_idx] + " iou: " + str(total_per_cat_iou[cat_idx])) + io.cprint("Max class iou:%.5f" % best_class_iou) state = { - 'model': model.module.state_dict() if torch.cuda.device_count() > 1 else model.state_dict(), - 'optimizer': opt.state_dict(), 'epoch': epoch, 'test_class_iou': best_class_iou} - torch.save(state, 'checkpoints/%s/best_clsiou_model.pth' % args.exp_name) + "model": model.module.state_dict() if torch.cuda.device_count() > 1 else model.state_dict(), + "optimizer": opt.state_dict(), + "epoch": epoch, + "test_class_iou": best_class_iou, + } + torch.save(state, "checkpoints/%s/best_clsiou_model.pth" % args.exp_name) # report best acc, ins_iou, cls_iou - io.cprint('Final Max Acc:%.5f' % best_acc) - io.cprint('Final Max instance iou:%.5f' % best_instance_iou) - io.cprint('Final Max class iou:%.5f' % best_class_iou) + io.cprint("Final Max Acc:%.5f" % best_acc) + io.cprint("Final Max instance iou:%.5f" % best_instance_iou) + io.cprint("Final Max class iou:%.5f" % best_class_iou) # save last model state = { - 'model': model.module.state_dict() if torch.cuda.device_count() > 1 else model.state_dict(), - 'optimizer': opt.state_dict(), 'epoch': args.epochs - 1, 'test_iou': best_instance_iou} - torch.save(state, 'checkpoints/%s/model_ep%d.pth' % (args.exp_name, args.epochs)) + "model": model.module.state_dict() if torch.cuda.device_count() > 1 else model.state_dict(), + "optimizer": opt.state_dict(), + "epoch": args.epochs - 1, + "test_iou": best_instance_iou, + } + torch.save(state, "checkpoints/%s/model_ep%d.pth" % (args.exp_name, args.epochs)) def train_epoch(train_loader, model, opt, scheduler, epoch, num_part, num_classes, io): @@ -175,22 +214,41 @@ def train_epoch(train_loader, model, opt, scheduler, epoch, num_part, num_classe metrics = defaultdict(lambda: list()) model.train() - for batch_id, (points, label, target, norm_plt) in tqdm(enumerate(train_loader), total=len(train_loader), smoothing=0.9): + for _batch_id, (points, label, target, norm_plt) in tqdm( + enumerate(train_loader), + total=len(train_loader), + smoothing=0.9, + ): batch_size, num_point, _ = points.size() - points, label, target, norm_plt = Variable(points.float()), Variable(label.long()), Variable(target.long()), \ - Variable(norm_plt.float()) + points, label, target, norm_plt = ( + Variable(points.float()), + Variable(label.long()), + Variable(target.long()), + Variable(norm_plt.float()), + ) points = points.transpose(2, 1) norm_plt = norm_plt.transpose(2, 1) - points, label, target, norm_plt = points.cuda(non_blocking=True), label.squeeze(1).cuda(non_blocking=True), \ - target.cuda(non_blocking=True), norm_plt.cuda(non_blocking=True) + points, label, target, norm_plt = ( + points.cuda(non_blocking=True), + label.squeeze(1).cuda(non_blocking=True), + target.cuda(non_blocking=True), + norm_plt.cuda(non_blocking=True), + ) # target: b,n seg_pred = model(points, norm_plt, to_categorical(label, num_classes)) # seg_pred: b,n,50 loss = F.nll_loss(seg_pred.contiguous().view(-1, num_part), target.view(-1, 1)[:, 0]) # instance iou without considering the class average at each batch_size: - batch_shapeious = compute_overall_iou(seg_pred, target, num_part) # list of of current batch_iou:[iou1,iou2,...,iou#b_size] + batch_shapeious = compute_overall_iou( + seg_pred, + target, + num_part, + ) # list of of current batch_iou:[iou1,iou2,...,iou#b_size] # total iou of current batch in each process: - batch_shapeious = seg_pred.new_tensor([np.sum(batch_shapeious)], dtype=torch.float64) # same device with seg_pred!!! + batch_shapeious = seg_pred.new_tensor( + [np.sum(batch_shapeious)], + dtype=torch.float64, + ) # same device with seg_pred!!! # Loss backward loss = torch.mean(loss) @@ -200,33 +258,37 @@ def train_epoch(train_loader, model, opt, scheduler, epoch, num_part, num_classe # accuracy seg_pred = seg_pred.contiguous().view(-1, num_part) # b*n,50 - target = target.view(-1, 1)[:, 0] # b*n + target = target.view(-1, 1)[:, 0] # b*n pred_choice = seg_pred.contiguous().data.max(1)[1] # b*n correct = pred_choice.eq(target.contiguous().data).sum() # torch.int64: total number of correct-predict pts # sum shape_ious += batch_shapeious.item() # count the sum of ious in each iteration - count += batch_size # count the total number of samples in each iteration + count += batch_size # count the total number of samples in each iteration train_loss += loss.item() * batch_size - accuracy.append(correct.item()/(batch_size * num_point)) # append the accuracy of each iteration + accuracy.append(correct.item() / (batch_size * num_point)) # append the accuracy of each iteration # Note: We do not need to calculate per_class iou during training - if args.scheduler == 'cos': + if args.scheduler == "cos": scheduler.step() - elif args.scheduler == 'step': - if opt.param_groups[0]['lr'] > 0.9e-5: + elif args.scheduler == "step": + if opt.param_groups[0]["lr"] > 0.9e-5: scheduler.step() - if opt.param_groups[0]['lr'] < 0.9e-5: + if opt.param_groups[0]["lr"] < 0.9e-5: for param_group in opt.param_groups: - param_group['lr'] = 0.9e-5 - io.cprint('Learning rate: %f' % opt.param_groups[0]['lr']) + param_group["lr"] = 0.9e-5 + io.cprint("Learning rate: %f" % opt.param_groups[0]["lr"]) - metrics['accuracy'] = np.mean(accuracy) - metrics['shape_avg_iou'] = shape_ious * 1.0 / count + metrics["accuracy"] = np.mean(accuracy) + metrics["shape_avg_iou"] = shape_ious * 1.0 / count - outstr = 'Train %d, loss: %f, train acc: %f, train ins_iou: %f' % (epoch+1, train_loss * 1.0 / count, - metrics['accuracy'], metrics['shape_avg_iou']) + outstr = "Train %d, loss: %f, train acc: %f, train ins_iou: %f" % ( + epoch + 1, + train_loss * 1.0 / count, + metrics["accuracy"], + metrics["shape_avg_iou"], + ) io.cprint(outstr) @@ -241,14 +303,26 @@ def test_epoch(test_loader, model, epoch, num_part, num_classes, io): model.eval() # label_size: b, means each sample has one corresponding class - for batch_id, (points, label, target, norm_plt) in tqdm(enumerate(test_loader), total=len(test_loader), smoothing=0.9): + for _batch_id, (points, label, target, norm_plt) in tqdm( + enumerate(test_loader), + total=len(test_loader), + smoothing=0.9, + ): batch_size, num_point, _ = points.size() - points, label, target, norm_plt = Variable(points.float()), Variable(label.long()), Variable(target.long()), \ - Variable(norm_plt.float()) + points, label, target, norm_plt = ( + Variable(points.float()), + Variable(label.long()), + Variable(target.long()), + Variable(norm_plt.float()), + ) points = points.transpose(2, 1) norm_plt = norm_plt.transpose(2, 1) - points, label, target, norm_plt = points.cuda(non_blocking=True), label.squeeze(1).cuda(non_blocking=True), \ - target.cuda(non_blocking=True), norm_plt.cuda(non_blocking=True) + points, label, target, norm_plt = ( + points.cuda(non_blocking=True), + label.squeeze(1).cuda(non_blocking=True), + target.cuda(non_blocking=True), + norm_plt.cuda(non_blocking=True), + ) seg_pred = model(points, norm_plt, to_categorical(label, num_classes)) # b,n,50 # instance iou without considering the class average at each batch_size: @@ -281,13 +355,19 @@ def test_epoch(test_loader, model, epoch, num_part, num_classes, io): for cat_idx in range(16): if final_total_per_cat_seen[cat_idx] > 0: # indicating this cat is included during previous iou appending - final_total_per_cat_iou[cat_idx] = final_total_per_cat_iou[cat_idx] / final_total_per_cat_seen[cat_idx] # avg class iou across all samples + final_total_per_cat_iou[cat_idx] = ( + final_total_per_cat_iou[cat_idx] / final_total_per_cat_seen[cat_idx] + ) # avg class iou across all samples - metrics['accuracy'] = np.mean(accuracy) - metrics['shape_avg_iou'] = shape_ious * 1.0 / count + metrics["accuracy"] = np.mean(accuracy) + metrics["shape_avg_iou"] = shape_ious * 1.0 / count - outstr = 'Test %d, loss: %f, test acc: %f test ins_iou: %f' % (epoch + 1, test_loss * 1.0 / count, - metrics['accuracy'], metrics['shape_avg_iou']) + outstr = "Test %d, loss: %f, test acc: %f test ins_iou: %f" % ( + epoch + 1, + test_loss * 1.0 / count, + metrics["accuracy"], + metrics["shape_avg_iou"], + ) io.cprint(outstr) @@ -296,11 +376,16 @@ def test_epoch(test_loader, model, epoch, num_part, num_classes, io): def test(args, io): # Dataloader - test_data = PartNormalDataset(npoints=2048, split='test', normalize=False) + test_data = PartNormalDataset(npoints=2048, split="test", normalize=False) print("The number of test data is:%d", len(test_data)) - test_loader = DataLoader(test_data, batch_size=args.test_batch_size, shuffle=False, num_workers=args.workers, - drop_last=False) + test_loader = DataLoader( + test_data, + batch_size=args.test_batch_size, + shuffle=False, + num_workers=args.workers, + drop_last=False, + ) # Try to load models num_part = 50 @@ -310,12 +395,15 @@ def test(args, io): io.cprint(str(model)) from collections import OrderedDict - state_dict = torch.load("checkpoints/%s/best_%s_model.pth" % (args.exp_name, args.model_type), - map_location=torch.device('cpu'))['model'] + + state_dict = torch.load( + f"checkpoints/{args.exp_name}/best_{args.model_type}_model.pth", + map_location=torch.device("cpu"), + )["model"] new_state_dict = OrderedDict() for layer in state_dict: - new_state_dict[layer.replace('module.', '')] = state_dict[layer] + new_state_dict[layer.replace("module.", "")] = state_dict[layer] model.load_state_dict(new_state_dict) model.eval() @@ -324,16 +412,29 @@ def test(args, io): metrics = defaultdict(lambda: list()) hist_acc = [] shape_ious = [] - total_per_cat_iou = np.zeros((16)).astype(np.float32) - total_per_cat_seen = np.zeros((16)).astype(np.int32) + total_per_cat_iou = np.zeros(16).astype(np.float32) + total_per_cat_seen = np.zeros(16).astype(np.int32) - for batch_id, (points, label, target, norm_plt) in tqdm(enumerate(test_loader), total=len(test_loader), smoothing=0.9): + for _batch_id, (points, label, target, norm_plt) in tqdm( + enumerate(test_loader), + total=len(test_loader), + smoothing=0.9, + ): batch_size, num_point, _ = points.size() - points, label, target, norm_plt = Variable(points.float()), Variable(label.long()), Variable(target.long()), Variable(norm_plt.float()) + points, label, target, norm_plt = ( + Variable(points.float()), + Variable(label.long()), + Variable(target.long()), + Variable(norm_plt.float()), + ) points = points.transpose(2, 1) norm_plt = norm_plt.transpose(2, 1) - points, label, target, norm_plt = points.cuda(non_blocking=True), label.squeeze().cuda( - non_blocking=True), target.cuda(non_blocking=True), norm_plt.cuda(non_blocking=True) + points, label, target, norm_plt = ( + points.cuda(non_blocking=True), + label.squeeze().cuda(non_blocking=True), + target.cuda(non_blocking=True), + norm_plt.cuda(non_blocking=True), + ) with torch.no_grad(): seg_pred = model(points, norm_plt, to_categorical(label, num_classes)) # b,n,50 @@ -353,11 +454,11 @@ def test(args, io): target = target.view(-1, 1)[:, 0] pred_choice = seg_pred.data.max(1)[1] correct = pred_choice.eq(target.data).cpu().sum() - metrics['accuracy'].append(correct.item() / (batch_size * num_point)) + metrics["accuracy"].append(correct.item() / (batch_size * num_point)) - hist_acc += metrics['accuracy'] - metrics['accuracy'] = np.mean(hist_acc) - metrics['shape_avg_iou'] = np.mean(shape_ious) + hist_acc += metrics["accuracy"] + metrics["accuracy"] = np.mean(hist_acc) + metrics["shape_avg_iou"] = np.mean(shape_ious) for cat_idx in range(16): if total_per_cat_seen[cat_idx] > 0: total_per_cat_iou[cat_idx] = total_per_cat_iou[cat_idx] / total_per_cat_seen[cat_idx] @@ -366,57 +467,51 @@ def test(args, io): class_iou = 0 for cat_idx in range(16): class_iou += total_per_cat_iou[cat_idx] - io.cprint(classes_str[cat_idx] + ' iou: ' + str(total_per_cat_iou[cat_idx])) # print the iou of each class + io.cprint(classes_str[cat_idx] + " iou: " + str(total_per_cat_iou[cat_idx])) # print the iou of each class avg_class_iou = class_iou / 16 - outstr = 'Test :: test acc: %f test class mIOU: %f, test instance mIOU: %f' % (metrics['accuracy'], avg_class_iou, metrics['shape_avg_iou']) + outstr = "Test :: test acc: {:f} test class mIOU: {:f}, test instance mIOU: {:f}".format( + metrics["accuracy"], + avg_class_iou, + metrics["shape_avg_iou"], + ) io.cprint(outstr) if __name__ == "__main__": # Training settings - parser = argparse.ArgumentParser(description='3D Shape Part Segmentation') - parser.add_argument('--model', type=str, default='PointMLP1') - parser.add_argument('--exp_name', type=str, default='demo1', metavar='N', - help='Name of the experiment') - parser.add_argument('--batch_size', type=int, default=32, metavar='batch_size', - help='Size of batch)') - parser.add_argument('--test_batch_size', type=int, default=32, metavar='batch_size', - help='Size of batch)') - parser.add_argument('--epochs', type=int, default=350, metavar='N', - help='number of episode to train') - parser.add_argument('--use_sgd', type=bool, default=False, - help='Use SGD') - parser.add_argument('--scheduler', type=str, default='step', - help='lr scheduler') - parser.add_argument('--step', type=int, default=40, - help='lr decay step') - parser.add_argument('--lr', type=float, default=0.003, metavar='LR', - help='learning rate') - parser.add_argument('--momentum', type=float, default=0.9, metavar='M', - help='SGD momentum (default: 0.9)') - parser.add_argument('--no_cuda', type=bool, default=False, - help='enables CUDA training') - parser.add_argument('--manual_seed', type=int, metavar='S', - help='random seed (default: 1)') - parser.add_argument('--eval', type=bool, default=False, - help='evaluate the model') - parser.add_argument('--num_points', type=int, default=2048, - help='num of points to use') - parser.add_argument('--workers', type=int, default=12) - parser.add_argument('--resume', type=bool, default=False, - help='Resume training or not') - parser.add_argument('--model_type', type=str, default='insiou', - help='choose to test the best insiou/clsiou/acc model (options: insiou, clsiou, acc)') + parser = argparse.ArgumentParser(description="3D Shape Part Segmentation") + parser.add_argument("--model", type=str, default="PointMLP1") + parser.add_argument("--exp_name", type=str, default="demo1", metavar="N", help="Name of the experiment") + parser.add_argument("--batch_size", type=int, default=32, metavar="batch_size", help="Size of batch)") + parser.add_argument("--test_batch_size", type=int, default=32, metavar="batch_size", help="Size of batch)") + parser.add_argument("--epochs", type=int, default=350, metavar="N", help="number of episode to train") + parser.add_argument("--use_sgd", type=bool, default=False, help="Use SGD") + parser.add_argument("--scheduler", type=str, default="step", help="lr scheduler") + parser.add_argument("--step", type=int, default=40, help="lr decay step") + parser.add_argument("--lr", type=float, default=0.003, metavar="LR", help="learning rate") + parser.add_argument("--momentum", type=float, default=0.9, metavar="M", help="SGD momentum (default: 0.9)") + parser.add_argument("--no_cuda", type=bool, default=False, help="enables CUDA training") + parser.add_argument("--manual_seed", type=int, metavar="S", help="random seed (default: 1)") + parser.add_argument("--eval", type=bool, default=False, help="evaluate the model") + parser.add_argument("--num_points", type=int, default=2048, help="num of points to use") + parser.add_argument("--workers", type=int, default=12) + parser.add_argument("--resume", type=bool, default=False, help="Resume training or not") + parser.add_argument( + "--model_type", + type=str, + default="insiou", + help="choose to test the best insiou/clsiou/acc model (options: insiou, clsiou, acc)", + ) args = parser.parse_args() - args.exp_name = args.model+"_"+args.exp_name + args.exp_name = args.model + "_" + args.exp_name _init_() if not args.eval: - io = IOStream('checkpoints/' + args.exp_name + '/%s_train.log' % (args.exp_name)) + io = IOStream("checkpoints/" + args.exp_name + "/%s_train.log" % (args.exp_name)) else: - io = IOStream('checkpoints/' + args.exp_name + '/%s_test.log' % (args.exp_name)) + io = IOStream("checkpoints/" + args.exp_name + "/%s_test.log" % (args.exp_name)) io.cprint(str(args)) if args.manual_seed is not None: @@ -427,12 +522,12 @@ if __name__ == "__main__": args.cuda = not args.no_cuda and torch.cuda.is_available() if args.cuda: - io.cprint('Using GPU') + io.cprint("Using GPU") if args.manual_seed is not None: torch.cuda.manual_seed(args.manual_seed) torch.cuda.manual_seed_all(args.manual_seed) else: - io.cprint('Using CPU') + io.cprint("Using CPU") if not args.eval: train(args, io) diff --git a/part_segmentation/model/__init__.py b/part_segmentation/model/__init__.py index 83a619e..159c4a2 100755 --- a/part_segmentation/model/__init__.py +++ b/part_segmentation/model/__init__.py @@ -1,2 +1 @@ -from __future__ import absolute_import from .pointMLP import pointMLP diff --git a/part_segmentation/model/pointMLP.py b/part_segmentation/model/pointMLP.py index 1b62414..76eee39 100644 --- a/part_segmentation/model/pointMLP.py +++ b/part_segmentation/model/pointMLP.py @@ -1,34 +1,32 @@ - import torch import torch.nn as nn import torch.nn.functional as F -from torch import einsum -from einops import rearrange, repeat + from pointnet2_ops import pointnet2_utils + def get_activation(activation): - if activation.lower() == 'gelu': + if activation.lower() == "gelu": return nn.GELU() - elif activation.lower() == 'rrelu': + elif activation.lower() == "rrelu": return nn.RReLU(inplace=True) - elif activation.lower() == 'selu': + elif activation.lower() == "selu": return nn.SELU(inplace=True) - elif activation.lower() == 'silu': + elif activation.lower() == "silu": return nn.SiLU(inplace=True) - elif activation.lower() == 'hardswish': + elif activation.lower() == "hardswish": return nn.Hardswish(inplace=True) - elif activation.lower() == 'leakyrelu': + elif activation.lower() == "leakyrelu": return nn.LeakyReLU(inplace=True) - elif activation.lower() == 'leakyrelu0.2': + elif activation.lower() == "leakyrelu0.2": return nn.LeakyReLU(negative_slope=0.2, inplace=True) else: return nn.ReLU(inplace=True) def square_distance(src, dst): - """ - Calculate Euclid distance between each two points. - src^T * dst = xn * xm + yn * ym + zn * zm； + """Calculate Euclid distance between each two points. + src^T * dst = xn * xm + yn * ym + zn * zm; sum(src^2, dim=-1) = xn*xn + yn*yn + zn*zn; sum(dst^2, dim=-1) = xm*xm + ym*ym + zm*zm; dist = (xn - xm)^2 + (yn - ym)^2 + (zn - zm)^2 @@ -37,23 +35,23 @@ def square_distance(src, dst): src: source points, [B, N, C] dst: target points, [B, M, C] Output: - dist: per-point square distance, [B, N, M] + dist: per-point square distance, [B, N, M]. """ B, N, _ = src.shape _, M, _ = dst.shape dist = -2 * torch.matmul(src, dst.permute(0, 2, 1)) - dist += torch.sum(src ** 2, -1).view(B, N, 1) - dist += torch.sum(dst ** 2, -1).view(B, 1, M) + dist += torch.sum(src**2, -1).view(B, N, 1) + dist += torch.sum(dst**2, -1).view(B, 1, M) return dist def index_points(points, idx): - """ - Input: + """Input: points: input points data, [B, N, C] - idx: sample index data, [B, S] + idx: sample index data, [B, S]. + Return: - new_points:, indexed points data, [B, S, C] + new_points:, indexed points data, [B, S, C]. """ device = points.device B = points.shape[0] @@ -62,17 +60,15 @@ def index_points(points, idx): repeat_shape = list(idx.shape) repeat_shape[0] = 1 batch_indices = torch.arange(B, dtype=torch.long).to(device).view(view_shape).repeat(repeat_shape) - new_points = points[batch_indices, idx, :] - return new_points + return points[batch_indices, idx, :] def farthest_point_sample(xyz, npoint): - """ - Input: + """Input: xyz: pointcloud data, [B, N, 3] npoint: number of samples Return: - centroids: sampled pointcloud index, [B, npoint] + centroids: sampled pointcloud index, [B, npoint]. """ device = xyz.device B, N, C = xyz.shape @@ -90,21 +86,21 @@ def farthest_point_sample(xyz, npoint): def query_ball_point(radius, nsample, xyz, new_xyz): - """ - Input: + """Input: radius: local region radius nsample: max sample number in local region xyz: all points, [B, N, 3] - new_xyz: query points, [B, S, 3] + new_xyz: query points, [B, S, 3]. + Return: - group_idx: grouped points index, [B, S, nsample] + group_idx: grouped points index, [B, S, nsample]. """ device = xyz.device B, N, C = xyz.shape _, S, _ = new_xyz.shape group_idx = torch.arange(N, dtype=torch.long).to(device).view(1, 1, N).repeat([B, S, 1]) sqrdists = square_distance(new_xyz, xyz) - group_idx[sqrdists > radius ** 2] = N + group_idx[sqrdists > radius**2] = N group_idx = group_idx.sort(dim=-1)[0][:, :, :nsample] group_first = group_idx[:, :, 0].view(B, S, 1).repeat([1, 1, nsample]) mask = group_idx == N @@ -113,13 +109,13 @@ def query_ball_point(radius, nsample, xyz, new_xyz): def knn_point(nsample, xyz, new_xyz): - """ - Input: + """Input: nsample: max sample number in local region xyz: all points, [B, N, C] - new_xyz: query points, [B, S, C] + new_xyz: query points, [B, S, C]. + Return: - group_idx: grouped points index, [B, S, nsample] + group_idx: grouped points index, [B, S, nsample]. """ sqrdists = square_distance(new_xyz, xyz) _, group_idx = torch.topk(sqrdists, nsample, dim=-1, largest=False, sorted=False) @@ -128,13 +124,12 @@ def knn_point(nsample, xyz, new_xyz): class LocalGrouper(nn.Module): def __init__(self, channel, groups, kneighbors, use_xyz=True, normalize="anchor", **kwargs): - """ - Give xyz[b,p,3] and fea[b,p,d], return new_xyz[b,g,3] and new_fea[b,g,k,d] + """Give xyz[b,p,3] and fea[b,p,d], return new_xyz[b,g,3] and new_fea[b,g,k,d] :param groups: groups number :param kneighbors: k-nerighbors - :param kwargs: others + :param kwargs: others. """ - super(LocalGrouper, self).__init__() + super().__init__() self.groups = groups self.kneighbors = kneighbors self.use_xyz = use_xyz @@ -143,11 +138,11 @@ class LocalGrouper(nn.Module): else: self.normalize = None if self.normalize not in ["center", "anchor"]: - print(f"Unrecognized normalize parameter (self.normalize), set to None. Should be one of [center, anchor].") + print("Unrecognized normalize parameter (self.normalize), set to None. Should be one of [center, anchor].") self.normalize = None if self.normalize is not None: - add_channel=3 if self.use_xyz else 0 - self.affine_alpha = nn.Parameter(torch.ones([1,1,1,channel + add_channel])) + add_channel = 3 if self.use_xyz else 0 + self.affine_alpha = nn.Parameter(torch.ones([1, 1, 1, channel + add_channel])) self.affine_beta = nn.Parameter(torch.zeros([1, 1, 1, channel + add_channel])) def forward(self, xyz, points): @@ -166,29 +161,33 @@ class LocalGrouper(nn.Module): grouped_xyz = index_points(xyz, idx) # [B, npoint, k, 3] grouped_points = index_points(points, idx) # [B, npoint, k, d] if self.use_xyz: - grouped_points = torch.cat([grouped_points, grouped_xyz],dim=-1) # [B, npoint, k, d+3] + grouped_points = torch.cat([grouped_points, grouped_xyz], dim=-1) # [B, npoint, k, d+3] if self.normalize is not None: - if self.normalize =="center": + if self.normalize == "center": mean = torch.mean(grouped_points, dim=2, keepdim=True) - if self.normalize =="anchor": - mean = torch.cat([new_points, new_xyz],dim=-1) if self.use_xyz else new_points - mean = mean.unsqueeze(dim=-2) # [B, npoint, 1, d+3] - std = torch.std((grouped_points-mean).reshape(B,-1),dim=-1,keepdim=True).unsqueeze(dim=-1).unsqueeze(dim=-1) - grouped_points = (grouped_points-mean)/(std + 1e-5) - grouped_points = self.affine_alpha*grouped_points + self.affine_beta + if self.normalize == "anchor": + mean = torch.cat([new_points, new_xyz], dim=-1) if self.use_xyz else new_points + mean = mean.unsqueeze(dim=-2) # [B, npoint, 1, d+3] + std = ( + torch.std((grouped_points - mean).reshape(B, -1), dim=-1, keepdim=True) + .unsqueeze(dim=-1) + .unsqueeze(dim=-1) + ) + grouped_points = (grouped_points - mean) / (std + 1e-5) + grouped_points = self.affine_alpha * grouped_points + self.affine_beta new_points = torch.cat([grouped_points, new_points.view(B, S, 1, -1).repeat(1, 1, self.kneighbors, 1)], dim=-1) return new_xyz, new_points class ConvBNReLU1D(nn.Module): - def __init__(self, in_channels, out_channels, kernel_size=1, bias=True, activation='relu'): - super(ConvBNReLU1D, self).__init__() + def __init__(self, in_channels, out_channels, kernel_size=1, bias=True, activation="relu"): + super().__init__() self.act = get_activation(activation) self.net = nn.Sequential( nn.Conv1d(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, bias=bias), nn.BatchNorm1d(out_channels), - self.act + self.act, ) def forward(self, x): @@ -196,30 +195,43 @@ class ConvBNReLU1D(nn.Module): class ConvBNReLURes1D(nn.Module): - def __init__(self, channel, kernel_size=1, groups=1, res_expansion=1.0, bias=True, activation='relu'): - super(ConvBNReLURes1D, self).__init__() + def __init__(self, channel, kernel_size=1, groups=1, res_expansion=1.0, bias=True, activation="relu"): + super().__init__() self.act = get_activation(activation) self.net1 = nn.Sequential( - nn.Conv1d(in_channels=channel, out_channels=int(channel * res_expansion), - kernel_size=kernel_size, groups=groups, bias=bias), + nn.Conv1d( + in_channels=channel, + out_channels=int(channel * res_expansion), + kernel_size=kernel_size, + groups=groups, + bias=bias, + ), nn.BatchNorm1d(int(channel * res_expansion)), - self.act + self.act, ) if groups > 1: self.net2 = nn.Sequential( - nn.Conv1d(in_channels=int(channel * res_expansion), out_channels=channel, - kernel_size=kernel_size, groups=groups, bias=bias), + nn.Conv1d( + in_channels=int(channel * res_expansion), + out_channels=channel, + kernel_size=kernel_size, + groups=groups, + bias=bias, + ), nn.BatchNorm1d(channel), self.act, - nn.Conv1d(in_channels=channel, out_channels=channel, - kernel_size=kernel_size, bias=bias), + nn.Conv1d(in_channels=channel, out_channels=channel, kernel_size=kernel_size, bias=bias), nn.BatchNorm1d(channel), ) else: self.net2 = nn.Sequential( - nn.Conv1d(in_channels=int(channel * res_expansion), out_channels=channel, - kernel_size=kernel_size, bias=bias), - nn.BatchNorm1d(channel) + nn.Conv1d( + in_channels=int(channel * res_expansion), + out_channels=channel, + kernel_size=kernel_size, + bias=bias, + ), + nn.BatchNorm1d(channel), ) def forward(self, x): @@ -227,21 +239,34 @@ class ConvBNReLURes1D(nn.Module): class PreExtraction(nn.Module): - def __init__(self, channels, out_channels, blocks=1, groups=1, res_expansion=1, bias=True, - activation='relu', use_xyz=True): - """ - input: [b,g,k,d]: output:[b,d,g] + def __init__( + self, + channels, + out_channels, + blocks=1, + groups=1, + res_expansion=1, + bias=True, + activation="relu", + use_xyz=True, + ): + """input: [b,g,k,d]: output:[b,d,g] :param channels: :param blocks: """ - super(PreExtraction, self).__init__() - in_channels = 3+2*channels if use_xyz else 2*channels + super().__init__() + in_channels = 3 + 2 * channels if use_xyz else 2 * channels self.transfer = ConvBNReLU1D(in_channels, out_channels, bias=bias, activation=activation) operation = [] for _ in range(blocks): operation.append( - ConvBNReLURes1D(out_channels, groups=groups, res_expansion=res_expansion, - bias=bias, activation=activation) + ConvBNReLURes1D( + out_channels, + groups=groups, + res_expansion=res_expansion, + bias=bias, + activation=activation, + ), ) self.operation = nn.Sequential(*operation) @@ -253,22 +278,20 @@ class PreExtraction(nn.Module): batch_size, _, _ = x.size() x = self.operation(x) # [b, d, k] x = F.adaptive_max_pool1d(x, 1).view(batch_size, -1) - x = x.reshape(b, n, -1).permute(0, 2, 1) - return x + return x.reshape(b, n, -1).permute(0, 2, 1) class PosExtraction(nn.Module): - def __init__(self, channels, blocks=1, groups=1, res_expansion=1, bias=True, activation='relu'): - """ - input[b,d,g]; output[b,d,g] + def __init__(self, channels, blocks=1, groups=1, res_expansion=1, bias=True, activation="relu"): + """input[b,d,g]; output[b,d,g] :param channels: :param blocks: """ - super(PosExtraction, self).__init__() + super().__init__() operation = [] for _ in range(blocks): operation.append( - ConvBNReLURes1D(channels, groups=groups, res_expansion=res_expansion, bias=bias, activation=activation) + ConvBNReLURes1D(channels, groups=groups, res_expansion=res_expansion, bias=bias, activation=activation), ) self.operation = nn.Sequential(*operation) @@ -277,22 +300,27 @@ class PosExtraction(nn.Module): class PointNetFeaturePropagation(nn.Module): - def __init__(self, in_channel, out_channel, blocks=1, groups=1, res_expansion=1.0, bias=True, activation='relu'): - super(PointNetFeaturePropagation, self).__init__() + def __init__(self, in_channel, out_channel, blocks=1, groups=1, res_expansion=1.0, bias=True, activation="relu"): + super().__init__() self.fuse = ConvBNReLU1D(in_channel, out_channel, 1, bias=bias) - self.extraction = PosExtraction(out_channel, blocks, groups=groups, - res_expansion=res_expansion, bias=bias, activation=activation) - + self.extraction = PosExtraction( + out_channel, + blocks, + groups=groups, + res_expansion=res_expansion, + bias=bias, + activation=activation, + ) def forward(self, xyz1, xyz2, points1, points2): - """ - Input: + """Input: xyz1: input points position data, [B, N, 3] xyz2: sampled input points position data, [B, S, 3] points1: input points data, [B, D', N] - points2: input points data, [B, D'', S] + points2: input points data, [B, D'', S]. + Return: - new_points: upsampled points data, [B, D''', N] + new_points: upsampled points data, [B, D''', N]. """ # xyz1 = xyz1.permute(0, 2, 1) # xyz2 = xyz2.permute(0, 2, 1) @@ -321,26 +349,40 @@ class PointNetFeaturePropagation(nn.Module): new_points = new_points.permute(0, 2, 1) new_points = self.fuse(new_points) - new_points = self.extraction(new_points) - return new_points - - + return self.extraction(new_points) class PointMLP(nn.Module): - def __init__(self, num_classes=50,points=2048, embed_dim=64, groups=1, res_expansion=1.0, - activation="relu", bias=True, use_xyz=True, normalize="anchor", - dim_expansion=[2, 2, 2, 2], pre_blocks=[2, 2, 2, 2], pos_blocks=[2, 2, 2, 2], - k_neighbors=[32, 32, 32, 32], reducers=[4, 4, 4, 4], - de_dims=[512, 256, 128, 128], de_blocks=[2,2,2,2], - gmp_dim=64,cls_dim=64, **kwargs): - super(PointMLP, self).__init__() + def __init__( + self, + num_classes=50, + points=2048, + embed_dim=64, + groups=1, + res_expansion=1.0, + activation="relu", + bias=True, + use_xyz=True, + normalize="anchor", + dim_expansion=[2, 2, 2, 2], + pre_blocks=[2, 2, 2, 2], + pos_blocks=[2, 2, 2, 2], + k_neighbors=[32, 32, 32, 32], + reducers=[4, 4, 4, 4], + de_dims=[512, 256, 128, 128], + de_blocks=[2, 2, 2, 2], + gmp_dim=64, + cls_dim=64, + **kwargs, + ): + super().__init__() self.stages = len(pre_blocks) self.class_num = num_classes self.points = points self.embedding = ConvBNReLU1D(6, embed_dim, bias=bias, activation=activation) - assert len(pre_blocks) == len(k_neighbors) == len(reducers) == len(pos_blocks) == len(dim_expansion), \ - "Please check stage number consistent for pre_blocks, pos_blocks k_neighbors, reducers." + assert ( + len(pre_blocks) == len(k_neighbors) == len(reducers) == len(pos_blocks) == len(dim_expansion) + ), "Please check stage number consistent for pre_blocks, pos_blocks k_neighbors, reducers." self.local_grouper_list = nn.ModuleList() self.pre_blocks_list = nn.ModuleList() self.pos_blocks_list = nn.ModuleList() @@ -359,29 +401,47 @@ class PointMLP(nn.Module): local_grouper = LocalGrouper(last_channel, anchor_points, kneighbor, use_xyz, normalize) # [b,g,k,d] self.local_grouper_list.append(local_grouper) # append pre_block_list - pre_block_module = PreExtraction(last_channel, out_channel, pre_block_num, groups=groups, - res_expansion=res_expansion, - bias=bias, activation=activation, use_xyz=use_xyz) + pre_block_module = PreExtraction( + last_channel, + out_channel, + pre_block_num, + groups=groups, + res_expansion=res_expansion, + bias=bias, + activation=activation, + use_xyz=use_xyz, + ) self.pre_blocks_list.append(pre_block_module) # append pos_block_list - pos_block_module = PosExtraction(out_channel, pos_block_num, groups=groups, - res_expansion=res_expansion, bias=bias, activation=activation) + pos_block_module = PosExtraction( + out_channel, + pos_block_num, + groups=groups, + res_expansion=res_expansion, + bias=bias, + activation=activation, + ) self.pos_blocks_list.append(pos_block_module) last_channel = out_channel en_dims.append(last_channel) - ### Building Decoder ##### self.decode_list = nn.ModuleList() en_dims.reverse() - de_dims.insert(0,en_dims[0]) - assert len(en_dims) ==len(de_dims) == len(de_blocks)+1 - for i in range(len(en_dims)-1): + de_dims.insert(0, en_dims[0]) + assert len(en_dims) == len(de_dims) == len(de_blocks) + 1 + for i in range(len(en_dims) - 1): self.decode_list.append( - PointNetFeaturePropagation(de_dims[i]+en_dims[i+1], de_dims[i+1], - blocks=de_blocks[i], groups=groups, res_expansion=res_expansion, - bias=bias, activation=activation) + PointNetFeaturePropagation( + de_dims[i] + en_dims[i + 1], + de_dims[i + 1], + blocks=de_blocks[i], + groups=groups, + res_expansion=res_expansion, + bias=bias, + activation=activation, + ), ) self.act = get_activation(activation) @@ -389,26 +449,26 @@ class PointMLP(nn.Module): # class label mapping self.cls_map = nn.Sequential( ConvBNReLU1D(16, cls_dim, bias=bias, activation=activation), - ConvBNReLU1D(cls_dim, cls_dim, bias=bias, activation=activation) + ConvBNReLU1D(cls_dim, cls_dim, bias=bias, activation=activation), ) # global max pooling mapping self.gmp_map_list = nn.ModuleList() for en_dim in en_dims: self.gmp_map_list.append(ConvBNReLU1D(en_dim, gmp_dim, bias=bias, activation=activation)) - self.gmp_map_end = ConvBNReLU1D(gmp_dim*len(en_dims), gmp_dim, bias=bias, activation=activation) + self.gmp_map_end = ConvBNReLU1D(gmp_dim * len(en_dims), gmp_dim, bias=bias, activation=activation) # classifier self.classifier = nn.Sequential( - nn.Conv1d(gmp_dim+cls_dim+de_dims[-1], 128, 1, bias=bias), + nn.Conv1d(gmp_dim + cls_dim + de_dims[-1], 128, 1, bias=bias), nn.BatchNorm1d(128), nn.Dropout(), - nn.Conv1d(128, num_classes, 1, bias=bias) + nn.Conv1d(128, num_classes, 1, bias=bias), ) self.en_dims = en_dims def forward(self, x, norm_plt, cls_label): xyz = x.permute(0, 2, 1) - x = torch.cat([x,norm_plt],dim=1) + x = torch.cat([x, norm_plt], dim=1) x = self.embedding(x) # B,D,N xyz_list = [xyz] # [B, N, 3] @@ -428,40 +488,54 @@ class PointMLP(nn.Module): x_list.reverse() x = x_list[0] for i in range(len(self.decode_list)): - x = self.decode_list[i](xyz_list[i+1], xyz_list[i], x_list[i+1],x) + x = self.decode_list[i](xyz_list[i + 1], xyz_list[i], x_list[i + 1], x) # here is the global context gmp_list = [] for i in range(len(x_list)): gmp_list.append(F.adaptive_max_pool1d(self.gmp_map_list[i](x_list[i]), 1)) - global_context = self.gmp_map_end(torch.cat(gmp_list, dim=1)) # [b, gmp_dim, 1] + global_context = self.gmp_map_end(torch.cat(gmp_list, dim=1)) # [b, gmp_dim, 1] - #here is the cls_token + # here is the cls_token cls_token = self.cls_map(cls_label.unsqueeze(dim=-1)) # [b, cls_dim, 1] x = torch.cat([x, global_context.repeat([1, 1, x.shape[-1]]), cls_token.repeat([1, 1, x.shape[-1]])], dim=1) x = self.classifier(x) x = F.log_softmax(x, dim=1) - x = x.permute(0, 2, 1) - return x + return x.permute(0, 2, 1) def pointMLP(num_classes=50, **kwargs) -> PointMLP: - return PointMLP(num_classes=num_classes, points=2048, embed_dim=64, groups=1, res_expansion=1.0, - activation="relu", bias=True, use_xyz=True, normalize="anchor", - dim_expansion=[2, 2, 2, 2], pre_blocks=[2, 2, 2, 2], pos_blocks=[2, 2, 2, 2], - k_neighbors=[32, 32, 32, 32], reducers=[4, 4, 4, 4], - de_dims=[512, 256, 128, 128], de_blocks=[4,4,4,4], - gmp_dim=64,cls_dim=64, **kwargs) + return PointMLP( + num_classes=num_classes, + points=2048, + embed_dim=64, + groups=1, + res_expansion=1.0, + activation="relu", + bias=True, + use_xyz=True, + normalize="anchor", + dim_expansion=[2, 2, 2, 2], + pre_blocks=[2, 2, 2, 2], + pos_blocks=[2, 2, 2, 2], + k_neighbors=[32, 32, 32, 32], + reducers=[4, 4, 4, 4], + de_dims=[512, 256, 128, 128], + de_blocks=[4, 4, 4, 4], + gmp_dim=64, + cls_dim=64, + **kwargs, + ) -if __name__ == '__main__': +if __name__ == "__main__": data = torch.rand(2, 3, 2048).cuda() norm = torch.rand(2, 3, 2048).cuda() cls_label = torch.rand([2, 16]).cuda() print(f"data shape: {data.shape}") print(f"norm shape: {norm.shape}") print(f"cls_label shape: {cls_label.shape}") - + print("===> testing pointMLP (segmentation) ...") model = pointMLP(50).cuda() out = model(data, norm, cls_label) # [2,2048,50] diff --git a/part_segmentation/util/data_util.py b/part_segmentation/util/data_util.py index 9725ab0..810324d 100755 --- a/part_segmentation/util/data_util.py +++ b/part_segmentation/util/data_util.py @@ -1,19 +1,21 @@ import glob +import json +import os + import h5py import numpy as np from torch.utils.data import Dataset -import os -import json + os.environ["HDF5_USE_FILE_LOCKING"] = "FALSE" def load_data(partition): all_data = [] all_label = [] - for h5_name in glob.glob('./data/modelnet40_ply_hdf5_2048/ply_data_%s*.h5' % partition): + for h5_name in glob.glob("./data/modelnet40_ply_hdf5_2048/ply_data_%s*.h5" % partition): f = h5py.File(h5_name) - data = f['data'][:].astype('float32') - label = f['label'][:].astype('int64') + data = f["data"][:].astype("float32") + label = f["label"][:].astype("int64") f.close() all_data.append(data) all_label.append(label) @@ -25,36 +27,34 @@ def load_data(partition): def pc_normalize(pc): centroid = np.mean(pc, axis=0) pc = pc - centroid - m = np.max(np.sqrt(np.sum(pc ** 2, axis=1))) - pc = pc / m - return pc + m = np.max(np.sqrt(np.sum(pc**2, axis=1))) + return pc / m def translate_pointcloud(pointcloud): - xyz1 = np.random.uniform(low=2./3., high=3./2., size=[3]) + xyz1 = np.random.uniform(low=2.0 / 3.0, high=3.0 / 2.0, size=[3]) xyz2 = np.random.uniform(low=-0.2, high=0.2, size=[3]) - - translated_pointcloud = np.add(np.multiply(pointcloud, xyz1), xyz2).astype('float32') - return translated_pointcloud + + return np.add(np.multiply(pointcloud, xyz1), xyz2).astype("float32") def jitter_pointcloud(pointcloud, sigma=0.01, clip=0.02): N, C = pointcloud.shape - pointcloud += np.clip(sigma * np.random.randn(N, C), -1*clip, clip) + pointcloud += np.clip(sigma * np.random.randn(N, C), -1 * clip, clip) return pointcloud # =========== ModelNet40 ================= class ModelNet40(Dataset): - def __init__(self, num_points, partition='train'): + def __init__(self, num_points, partition="train"): self.data, self.label = load_data(partition) self.num_points = num_points self.partition = partition # Here the new given partition will cover the 'train' def __getitem__(self, item): # indice of the pts or label - pointcloud = self.data[item][:self.num_points] + pointcloud = self.data[item][: self.num_points] label = self.label[item] - if self.partition == 'train': + if self.partition == "train": # pointcloud = pc_normalize(pointcloud) # you can try to add it or not to train our model pointcloud = translate_pointcloud(pointcloud) np.random.shuffle(pointcloud) # shuffle the order of pts @@ -66,46 +66,46 @@ class ModelNet40(Dataset): # =========== ShapeNet Part ================= class PartNormalDataset(Dataset): - def __init__(self, npoints=2500, split='train', normalize=False): + def __init__(self, npoints=2500, split="train", normalize=False): self.npoints = npoints - self.root = './data/shapenetcore_partanno_segmentation_benchmark_v0_normal' - self.catfile = os.path.join(self.root, 'synsetoffset2category.txt') + self.root = "./data/shapenetcore_partanno_segmentation_benchmark_v0_normal" + self.catfile = os.path.join(self.root, "synsetoffset2category.txt") self.cat = {} self.normalize = normalize - with open(self.catfile, 'r') as f: + with open(self.catfile) as f: for line in f: ls = line.strip().split() self.cat[ls[0]] = ls[1] self.cat = {k: v for k, v in self.cat.items()} self.meta = {} - with open(os.path.join(self.root, 'train_test_split', 'shuffled_train_file_list.json'), 'r') as f: - train_ids = set([str(d.split('/')[2]) for d in json.load(f)]) - with open(os.path.join(self.root, 'train_test_split', 'shuffled_val_file_list.json'), 'r') as f: - val_ids = set([str(d.split('/')[2]) for d in json.load(f)]) - with open(os.path.join(self.root, 'train_test_split', 'shuffled_test_file_list.json'), 'r') as f: - test_ids = set([str(d.split('/')[2]) for d in json.load(f)]) + with open(os.path.join(self.root, "train_test_split", "shuffled_train_file_list.json")) as f: + train_ids = set([str(d.split("/")[2]) for d in json.load(f)]) + with open(os.path.join(self.root, "train_test_split", "shuffled_val_file_list.json")) as f: + val_ids = set([str(d.split("/")[2]) for d in json.load(f)]) + with open(os.path.join(self.root, "train_test_split", "shuffled_test_file_list.json")) as f: + test_ids = set([str(d.split("/")[2]) for d in json.load(f)]) for item in self.cat: self.meta[item] = [] dir_point = os.path.join(self.root, self.cat[item]) fns = sorted(os.listdir(dir_point)) - if split == 'trainval': + if split == "trainval": fns = [fn for fn in fns if ((fn[0:-4] in train_ids) or (fn[0:-4] in val_ids))] - elif split == 'train': + elif split == "train": fns = [fn for fn in fns if fn[0:-4] in train_ids] - elif split == 'val': + elif split == "val": fns = [fn for fn in fns if fn[0:-4] in val_ids] - elif split == 'test': + elif split == "test": fns = [fn for fn in fns if fn[0:-4] in test_ids] else: - print('Unknown split: %s. Exiting..' % (split)) + print("Unknown split: %s. Exiting.." % (split)) exit(-1) for fn in fns: - token = (os.path.splitext(os.path.basename(fn))[0]) - self.meta[item].append(os.path.join(dir_point, token + '.txt')) + token = os.path.splitext(os.path.basename(fn))[0] + self.meta[item].append(os.path.join(dir_point, token + ".txt")) self.datapath = [] for item in self.cat: @@ -114,11 +114,24 @@ class PartNormalDataset(Dataset): self.classes = dict(zip(self.cat, range(len(self.cat)))) # Mapping from category ('Chair') to a list of int [10,11,12,13] as segmentation labels - self.seg_classes = {'Earphone': [16, 17, 18], 'Motorbike': [30, 31, 32, 33, 34, 35], 'Rocket': [41, 42, 43], - 'Car': [8, 9, 10, 11], 'Laptop': [28, 29], 'Cap': [6, 7], 'Skateboard': [44, 45, 46], - 'Mug': [36, 37], 'Guitar': [19, 20, 21], 'Bag': [4, 5], 'Lamp': [24, 25, 26, 27], - 'Table': [47, 48, 49], 'Airplane': [0, 1, 2, 3], 'Pistol': [38, 39, 40], - 'Chair': [12, 13, 14, 15], 'Knife': [22, 23]} + self.seg_classes = { + "Earphone": [16, 17, 18], + "Motorbike": [30, 31, 32, 33, 34, 35], + "Rocket": [41, 42, 43], + "Car": [8, 9, 10, 11], + "Laptop": [28, 29], + "Cap": [6, 7], + "Skateboard": [44, 45, 46], + "Mug": [36, 37], + "Guitar": [19, 20, 21], + "Bag": [4, 5], + "Lamp": [24, 25, 26, 27], + "Table": [47, 48, 49], + "Airplane": [0, 1, 2, 3], + "Pistol": [38, 39, 40], + "Chair": [12, 13, 14, 15], + "Knife": [22, 23], + } self.cache = {} # from index to (point_set, cls, seg) tuple self.cache_size = 20000 @@ -156,9 +169,9 @@ class PartNormalDataset(Dataset): return len(self.datapath) -if __name__ == '__main__': - train = PartNormalDataset(npoints=2048, split='trainval', normalize=False) - test = PartNormalDataset(npoints=2048, split='test', normalize=False) +if __name__ == "__main__": + train = PartNormalDataset(npoints=2048, split="trainval", normalize=False) + test = PartNormalDataset(npoints=2048, split="test", normalize=False) for data, label, _, _ in train: print(data.shape) print(label.shape) diff --git a/part_segmentation/util/util.py b/part_segmentation/util/util.py index 00afdd8..c3a6463 100755 --- a/part_segmentation/util/util.py +++ b/part_segmentation/util/util.py @@ -4,9 +4,8 @@ import torch.nn.functional as F def cal_loss(pred, gold, smoothing=True): - ''' Calculate cross entropy loss, apply label smoothing if needed. ''' - - gold = gold.contiguous().view(-1) # gold is the groudtruth label in the dataloader + """Calculate cross entropy loss, apply label smoothing if needed.""" + gold = gold.contiguous().view(-1) # gold is the groudtruth label in the dataloader if smoothing: eps = 0.2 @@ -18,19 +17,19 @@ def cal_loss(pred, gold, smoothing=True): loss = -(one_hot * log_prb).sum(dim=1).mean() else: - loss = F.cross_entropy(pred, gold, reduction='mean') + loss = F.cross_entropy(pred, gold, reduction="mean") return loss # create a file and write the text into it: -class IOStream(): +class IOStream: def __init__(self, path): - self.f = open(path, 'a') + self.f = open(path, "a") def cprint(self, text): print(text) - self.f.write(text+'\n') + self.f.write(text + "\n") self.f.flush() def close(self): @@ -38,22 +37,24 @@ class IOStream(): def to_categorical(y, num_classes): - """ 1-hot encodes a tensor """ + """1-hot encodes a tensor.""" new_y = torch.eye(num_classes)[y.cpu().data.numpy(),] - if (y.is_cuda): + if y.is_cuda: return new_y.cuda(non_blocking=True) return new_y def compute_overall_iou(pred, target, num_classes): shape_ious = [] - pred = pred.max(dim=2)[1] # (batch_size, num_points) the pred_class_idx of each point in each sample + pred = pred.max(dim=2)[1] # (batch_size, num_points) the pred_class_idx of each point in each sample pred_np = pred.cpu().data.numpy() target_np = target.cpu().data.numpy() - for shape_idx in range(pred.size(0)): # sample_idx + for shape_idx in range(pred.size(0)): # sample_idx part_ious = [] - for part in range(num_classes): # class_idx! no matter which category, only consider all part_classes of all categories, check all 50 classes + for part in range( + num_classes, + ): # class_idx! no matter which category, only consider all part_classes of all categories, check all 50 classes # for target, each point has a class no matter which category owns this point! also 50 classes!!! # only return 1 when both belongs to this class, which means correct: I = np.sum(np.logical_and(pred_np[shape_idx] == part, target_np[shape_idx] == part)) @@ -63,7 +64,9 @@ def compute_overall_iou(pred, target, num_classes): F = np.sum(target_np[shape_idx] == part) if F != 0: - iou = I / float(U) # iou across all points for this class - part_ious.append(iou) # append the iou of this class - shape_ious.append(np.mean(part_ious)) # each time append an average iou across all classes of this sample (sample_level!) - return shape_ious # [batch_size] + iou = I / float(U) # iou across all points for this class + part_ious.append(iou) # append the iou of this class + shape_ious.append( + np.mean(part_ious), + ) # each time append an average iou across all classes of this sample (sample_level!) + return shape_ious # [batch_size] diff --git a/pointnet2_ops_lib/pointnet2_ops/pointnet2_modules.py b/pointnet2_ops_lib/pointnet2_ops/pointnet2_modules.py index a0ad4f6..003e1db 100644 --- a/pointnet2_ops_lib/pointnet2_ops/pointnet2_modules.py +++ b/pointnet2_ops_lib/pointnet2_ops/pointnet2_modules.py @@ -1,16 +1,15 @@ -from typing import List, Optional, Tuple - import torch import torch.nn as nn import torch.nn.functional as F + from pointnet2_ops import pointnet2_utils -def build_shared_mlp(mlp_spec: List[int], bn: bool = True): +def build_shared_mlp(mlp_spec: list[int], bn: bool = True): layers = [] for i in range(1, len(mlp_spec)): layers.append( - nn.Conv2d(mlp_spec[i - 1], mlp_spec[i], kernel_size=1, bias=not bn) + nn.Conv2d(mlp_spec[i - 1], mlp_spec[i], kernel_size=1, bias=not bn), ) if bn: layers.append(nn.BatchNorm2d(mlp_spec[i])) @@ -21,36 +20,37 @@ def build_shared_mlp(mlp_spec: List[int], bn: bool = True): class _PointnetSAModuleBase(nn.Module): def __init__(self): - super(_PointnetSAModuleBase, self).__init__() + super().__init__() self.npoint = None self.groupers = None self.mlps = None def forward( - self, xyz: torch.Tensor, features: Optional[torch.Tensor] - ) -> Tuple[torch.Tensor, torch.Tensor]: - r""" - Parameters + self, + xyz: torch.Tensor, + features: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor]: + r"""Parameters ---------- xyz : torch.Tensor (B, N, 3) tensor of the xyz coordinates of the features features : torch.Tensor (B, C, N) tensor of the descriptors of the the features - Returns + Returns: ------- new_xyz : torch.Tensor (B, npoint, 3) tensor of the new features' xyz new_features : torch.Tensor (B, \sum_k(mlps[k][-1]), npoint) tensor of the new_features descriptors """ - new_features_list = [] xyz_flipped = xyz.transpose(1, 2).contiguous() new_xyz = ( pointnet2_utils.gather_operation( - xyz_flipped, pointnet2_utils.furthest_point_sample(xyz, self.npoint) + xyz_flipped, + pointnet2_utils.furthest_point_sample(xyz, self.npoint), ) .transpose(1, 2) .contiguous() @@ -60,12 +60,15 @@ class _PointnetSAModuleBase(nn.Module): for i in range(len(self.groupers)): new_features = self.groupers[i]( - xyz, new_xyz, features + xyz, + new_xyz, + features, ) # (B, C, npoint, nsample) new_features = self.mlps[i](new_features) # (B, mlp[-1], npoint, nsample) new_features = F.max_pool2d( - new_features, kernel_size=[1, new_features.size(3)] + new_features, + kernel_size=[1, new_features.size(3)], ) # (B, mlp[-1], npoint, 1) new_features = new_features.squeeze(-1) # (B, mlp[-1], npoint) @@ -75,7 +78,7 @@ class _PointnetSAModuleBase(nn.Module): class PointnetSAModuleMSG(_PointnetSAModuleBase): - r"""Pointnet set abstrction layer with multiscale grouping + r"""Pointnet set abstrction layer with multiscale grouping. Parameters ---------- @@ -93,7 +96,7 @@ class PointnetSAModuleMSG(_PointnetSAModuleBase): def __init__(self, npoint, radii, nsamples, mlps, bn=True, use_xyz=True): # type: (PointnetSAModuleMSG, int, List[float], List[int], List[List[int]], bool, bool) -> None - super(PointnetSAModuleMSG, self).__init__() + super().__init__() assert len(radii) == len(nsamples) == len(mlps) @@ -106,7 +109,7 @@ class PointnetSAModuleMSG(_PointnetSAModuleBase): self.groupers.append( pointnet2_utils.QueryAndGroup(radius, nsample, use_xyz=use_xyz) if npoint is not None - else pointnet2_utils.GroupAll(use_xyz) + else pointnet2_utils.GroupAll(use_xyz), ) mlp_spec = mlps[i] if use_xyz: @@ -116,7 +119,7 @@ class PointnetSAModuleMSG(_PointnetSAModuleBase): class PointnetSAModule(PointnetSAModuleMSG): - r"""Pointnet set abstrction layer + r"""Pointnet set abstrction layer. Parameters ---------- @@ -133,10 +136,16 @@ class PointnetSAModule(PointnetSAModuleMSG): """ def __init__( - self, mlp, npoint=None, radius=None, nsample=None, bn=True, use_xyz=True + self, + mlp, + npoint=None, + radius=None, + nsample=None, + bn=True, + use_xyz=True, ): # type: (PointnetSAModule, List[int], int, float, int, bool, bool) -> None - super(PointnetSAModule, self).__init__( + super().__init__( mlps=[mlp], npoint=npoint, radii=[radius], @@ -147,7 +156,7 @@ class PointnetSAModule(PointnetSAModuleMSG): class PointnetFPModule(nn.Module): - r"""Propigates the features of one set to another + r"""Propigates the features of one set to another. Parameters ---------- @@ -159,13 +168,12 @@ class PointnetFPModule(nn.Module): def __init__(self, mlp, bn=True): # type: (PointnetFPModule, List[int], bool) -> None - super(PointnetFPModule, self).__init__() + super().__init__() self.mlp = build_shared_mlp(mlp, bn=bn) def forward(self, unknown, known, unknow_feats, known_feats): # type: (PointnetFPModule, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor) -> torch.Tensor - r""" - Parameters + r"""Parameters ---------- unknown : torch.Tensor (B, n, 3) tensor of the xyz positions of the unknown features @@ -176,12 +184,11 @@ class PointnetFPModule(nn.Module): known_feats : torch.Tensor (B, C2, m) tensor of features to be propigated - Returns + Returns: ------- new_features : torch.Tensor (B, mlp[-1], n) tensor of the features of the unknown features """ - if known is not None: dist, idx = pointnet2_utils.three_nn(unknown, known) dist_recip = 1.0 / (dist + 1e-8) @@ -189,16 +196,19 @@ class PointnetFPModule(nn.Module): weight = dist_recip / norm interpolated_feats = pointnet2_utils.three_interpolate( - known_feats, idx, weight + known_feats, + idx, + weight, ) else: interpolated_feats = known_feats.expand( - *(known_feats.size()[0:2] + [unknown.size(1)]) + *(known_feats.size()[0:2] + [unknown.size(1)]), ) if unknow_feats is not None: new_features = torch.cat( - [interpolated_feats, unknow_feats], dim=1 + [interpolated_feats, unknow_feats], + dim=1, ) # (B, C2 + C1, n) else: new_features = interpolated_feats diff --git a/pointnet2_ops_lib/pointnet2_ops/pointnet2_utils.py b/pointnet2_ops_lib/pointnet2_ops/pointnet2_utils.py index 150fccc..4f2b789 100644 --- a/pointnet2_ops_lib/pointnet2_ops/pointnet2_utils.py +++ b/pointnet2_ops_lib/pointnet2_ops/pointnet2_utils.py @@ -1,22 +1,24 @@ +import warnings +from typing import * + import torch import torch.nn as nn -import warnings from torch.autograd import Function -from typing import * try: import pointnet2_ops._ext as _ext except ImportError: - from torch.utils.cpp_extension import load import glob - import os.path as osp import os + import os.path as osp + + from torch.utils.cpp_extension import load warnings.warn("Unable to load pointnet2_ops cpp extension. JIT Compiling.") _ext_src_root = osp.join(osp.dirname(__file__), "_ext-src") _ext_sources = glob.glob(osp.join(_ext_src_root, "src", "*.cpp")) + glob.glob( - osp.join(_ext_src_root, "src", "*.cu") + osp.join(_ext_src_root, "src", "*.cu"), ) _ext_headers = glob.glob(osp.join(_ext_src_root, "include", "*")) @@ -35,9 +37,8 @@ class FurthestPointSampling(Function): @staticmethod def forward(ctx, xyz, npoint): # type: (Any, torch.Tensor, int) -> torch.Tensor - r""" - Uses iterative furthest point sampling to select a set of npoint features that have the largest - minimum distance + r"""Uses iterative furthest point sampling to select a set of npoint features that have the largest + minimum distance. Parameters ---------- @@ -46,7 +47,7 @@ class FurthestPointSampling(Function): npoint : int32 number of features in the sampled set - Returns + Returns: ------- torch.Tensor (B, npoint) tensor containing the set @@ -69,9 +70,7 @@ class GatherOperation(Function): @staticmethod def forward(ctx, features, idx): # type: (Any, torch.Tensor, torch.Tensor) -> torch.Tensor - r""" - - Parameters + r"""Parameters ---------- features : torch.Tensor (B, C, N) tensor @@ -79,12 +78,11 @@ class GatherOperation(Function): idx : torch.Tensor (B, npoint) tensor of the features to gather - Returns + Returns: ------- torch.Tensor (B, C, npoint) tensor """ - ctx.save_for_backward(idx, features) return _ext.gather_points(features, idx) @@ -105,16 +103,15 @@ class ThreeNN(Function): @staticmethod def forward(ctx, unknown, known): # type: (Any, torch.Tensor, torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor] - r""" - Find the three nearest neighbors of unknown in known + r"""Find the three nearest neighbors of unknown in known Parameters ---------- unknown : torch.Tensor (B, n, 3) tensor of known features known : torch.Tensor - (B, m, 3) tensor of unknown features + (B, m, 3) tensor of unknown features. - Returns + Returns: ------- dist : torch.Tensor (B, n, 3) l2 distance to the three nearest neighbors @@ -140,8 +137,7 @@ class ThreeInterpolate(Function): @staticmethod def forward(ctx, features, idx, weight): # type(Any, torch.Tensor, torch.Tensor, torch.Tensor) -> Torch.Tensor - r""" - Performs weight linear interpolation on 3 features + r"""Performs weight linear interpolation on 3 features Parameters ---------- features : torch.Tensor @@ -149,9 +145,9 @@ class ThreeInterpolate(Function): idx : torch.Tensor (B, n, 3) three nearest neighbors of the target features in features weight : torch.Tensor - (B, n, 3) weights + (B, n, 3) weights. - Returns + Returns: ------- torch.Tensor (B, c, n) tensor of the interpolated features @@ -163,13 +159,12 @@ class ThreeInterpolate(Function): @staticmethod def backward(ctx, grad_out): # type: (Any, torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor] - r""" - Parameters + r"""Parameters ---------- grad_out : torch.Tensor (B, c, n) tensor with gradients of ouputs - Returns + Returns: ------- grad_features : torch.Tensor (B, c, m) tensor with gradients of features @@ -182,7 +177,10 @@ class ThreeInterpolate(Function): m = features.size(2) grad_features = _ext.three_interpolate_grad( - grad_out.contiguous(), idx, weight, m + grad_out.contiguous(), + idx, + weight, + m, ) return grad_features, torch.zeros_like(idx), torch.zeros_like(weight) @@ -195,16 +193,14 @@ class GroupingOperation(Function): @staticmethod def forward(ctx, features, idx): # type: (Any, torch.Tensor, torch.Tensor) -> torch.Tensor - r""" - - Parameters + r"""Parameters ---------- features : torch.Tensor (B, C, N) tensor of features to group idx : torch.Tensor (B, npoint, nsample) tensor containing the indicies of features to group with - Returns + Returns: ------- torch.Tensor (B, C, npoint, nsample) tensor @@ -216,14 +212,12 @@ class GroupingOperation(Function): @staticmethod def backward(ctx, grad_out): # type: (Any, torch.tensor) -> Tuple[torch.Tensor, torch.Tensor] - r""" - - Parameters + r"""Parameters ---------- grad_out : torch.Tensor (B, C, npoint, nsample) tensor of the gradients of the output from forward - Returns + Returns: ------- torch.Tensor (B, C, N) gradient of the features @@ -244,9 +238,7 @@ class BallQuery(Function): @staticmethod def forward(ctx, radius, nsample, xyz, new_xyz): # type: (Any, float, int, torch.Tensor, torch.Tensor) -> torch.Tensor - r""" - - Parameters + r"""Parameters ---------- radius : float radius of the balls @@ -257,7 +249,7 @@ class BallQuery(Function): new_xyz : torch.Tensor (B, npoint, 3) centers of the ball query - Returns + Returns: ------- torch.Tensor (B, npoint, nsample) tensor with the indicies of the features that form the query balls @@ -277,8 +269,7 @@ ball_query = BallQuery.apply class QueryAndGroup(nn.Module): - r""" - Groups with a ball query of radius + r"""Groups with a ball query of radius. Parameters --------- @@ -290,13 +281,12 @@ class QueryAndGroup(nn.Module): def __init__(self, radius, nsample, use_xyz=True): # type: (QueryAndGroup, float, int, bool) -> None - super(QueryAndGroup, self).__init__() + super().__init__() self.radius, self.nsample, self.use_xyz = radius, nsample, use_xyz def forward(self, xyz, new_xyz, features=None): # type: (QueryAndGroup, torch.Tensor. torch.Tensor, torch.Tensor) -> Tuple[Torch.Tensor] - r""" - Parameters + r"""Parameters ---------- xyz : torch.Tensor xyz coordinates of the features (B, N, 3) @@ -305,12 +295,11 @@ class QueryAndGroup(nn.Module): features : torch.Tensor Descriptors of the features (B, C, N) - Returns + Returns: ------- new_features : torch.Tensor (B, 3 + C, npoint, nsample) tensor """ - idx = ball_query(self.radius, self.nsample, xyz, new_xyz) xyz_trans = xyz.transpose(1, 2).contiguous() grouped_xyz = grouping_operation(xyz_trans, idx) # (B, 3, npoint, nsample) @@ -320,22 +309,20 @@ class QueryAndGroup(nn.Module): grouped_features = grouping_operation(features, idx) if self.use_xyz: new_features = torch.cat( - [grouped_xyz, grouped_features], dim=1 + [grouped_xyz, grouped_features], + dim=1, ) # (B, C + 3, npoint, nsample) else: new_features = grouped_features else: - assert ( - self.use_xyz - ), "Cannot have not features and not use xyz as a feature!" + assert self.use_xyz, "Cannot have not features and not use xyz as a feature!" new_features = grouped_xyz return new_features class GroupAll(nn.Module): - r""" - Groups all features + r"""Groups all features. Parameters --------- @@ -343,13 +330,12 @@ class GroupAll(nn.Module): def __init__(self, use_xyz=True): # type: (GroupAll, bool) -> None - super(GroupAll, self).__init__() + super().__init__() self.use_xyz = use_xyz def forward(self, xyz, new_xyz, features=None): # type: (GroupAll, torch.Tensor, torch.Tensor, torch.Tensor) -> Tuple[torch.Tensor] - r""" - Parameters + r"""Parameters ---------- xyz : torch.Tensor xyz coordinates of the features (B, N, 3) @@ -358,18 +344,18 @@ class GroupAll(nn.Module): features : torch.Tensor Descriptors of the features (B, C, N) - Returns + Returns: ------- new_features : torch.Tensor (B, C + 3, 1, N) tensor """ - grouped_xyz = xyz.transpose(1, 2).unsqueeze(2) if features is not None: grouped_features = features.unsqueeze(2) if self.use_xyz: new_features = torch.cat( - [grouped_xyz, grouped_features], dim=1 + [grouped_xyz, grouped_features], + dim=1, ) # (B, 3 + C, 1, N) else: new_features = grouped_features diff --git a/pointnet2_ops_lib/setup.py b/pointnet2_ops_lib/setup.py index faf7154..822fbe1 100644 --- a/pointnet2_ops_lib/setup.py +++ b/pointnet2_ops_lib/setup.py @@ -8,7 +8,7 @@ from torch.utils.cpp_extension import BuildExtension, CUDAExtension this_dir = osp.dirname(osp.abspath(__file__)) _ext_src_root = osp.join("pointnet2_ops", "_ext-src") _ext_sources = glob.glob(osp.join(_ext_src_root, "src", "*.cpp")) + glob.glob( - osp.join(_ext_src_root, "src", "*.cu") + osp.join(_ext_src_root, "src", "*.cu"), ) _ext_headers = glob.glob(osp.join(_ext_src_root, "include", "*")) @@ -32,7 +32,7 @@ setup( "nvcc": ["-O3", "-Xfatbin", "-compress-all"], }, include_dirs=[osp.join(this_dir, _ext_src_root, "include")], - ) + ), ], cmdclass={"build_ext": BuildExtension}, include_package_data=True,