Safran/PointMLP
1
0
Fork 0
Code Issues Pull requests Packages Projects Releases Wiki Activity

auto-formatting

Browse source
This commit is contained in:
Laurent FAINSIN 2023-08-03 16:40:14 +02:00
parent cc00fa7215
commit f680d6f7a8
40 changed files with 1822 additions and 1385 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

11
analysis.py
View file

@ -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")

76
classification_ModelNet40/data.py
View file

@ -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__()}")

6
classification_ModelNet40/helper.py
View file

@ -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

197
classification_ModelNet40/main.py
View file

@ -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()

2
classification_ModelNet40/models/__init__.py
View file

@ -1,3 +1 @@
from __future__ import absolute_import
from .pointmlp import pointMLP, pointMLPElite

280
classification_ModelNet40/models/pointmlp.py
View file

@ -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)

79
classification_ModelNet40/test.py
View file

@ -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()

5
classification_ModelNet40/utils/__init__.py
View file

@ -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

86
classification_ModelNet40/utils/logger.py
View file

@ -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')
savefig("test.eps")

152
classification_ModelNet40/utils/misc.py
View file

@ -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

14
classification_ModelNet40/utils/progress/progress/__init__.py
View file

@ -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):

51
classification_ModelNet40/utils/progress/progress/bar.py
View file

@ -1,5 +1,3 @@
# -*- coding: utf-8 -*-
# Copyright (c) 2012 Giorgos Verigakis <verigak@gmail.com>
#
# 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 = (" ", "░", "▒", "▓", "█")

9
classification_ModelNet40/utils/progress/progress/counter.py
View file

@ -1,5 +1,3 @@
# -*- coding: utf-8 -*-
# Copyright (c) 2012 Giorgos Verigakis <verigak@gmail.com>
#
# 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 = ("○", "◔", "◑", "◕", "●")

42
classification_ModelNet40/utils/progress/progress/helpers.py
View file

@ -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):

16
classification_ModelNet40/utils/progress/progress/spinner.py
View file

@ -1,5 +1,3 @@
# -*- coding: utf-8 -*-
# Copyright (c) 2012 Giorgos Verigakis <verigak@gmail.com>
#
# 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 = ["⣾", "⣷", "⣯", "⣟", "⡿", "⢿", "⣻", "⣽"]

40
classification_ModelNet40/utils/progress/setup.py
View file

@ -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",
],
)

26
classification_ModelNet40/utils/progress/test_progress.py
View file

@ -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()

132
classification_ModelNet40/voting.py
View file

@ -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()

40
classification_ScanObjectNN/ScanObjectNN.py
View file

@ -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)

199
classification_ScanObjectNN/main.py
View file

@ -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()

2
classification_ScanObjectNN/models/__init__.py
View file

@ -1,3 +1 @@
from __future__ import absolute_import
from .pointmlp import pointMLP, pointMLPElite

281
classification_ScanObjectNN/models/pointmlp.py
View file

@ -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)

5
classification_ScanObjectNN/utils/__init__.py
View file

@ -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

86
classification_ScanObjectNN/utils/logger.py
View file

@ -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')
savefig("test.eps")

152
classification_ScanObjectNN/utils/misc.py
View file

@ -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

14
classification_ScanObjectNN/utils/progress/progress/__init__.py
View file

@ -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):

51
classification_ScanObjectNN/utils/progress/progress/bar.py
View file

@ -1,5 +1,3 @@
# -*- coding: utf-8 -*-
# Copyright (c) 2012 Giorgos Verigakis <verigak@gmail.com>
#
# 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 = (" ", "░", "▒", "▓", "█")

9
classification_ScanObjectNN/utils/progress/progress/counter.py
View file

@ -1,5 +1,3 @@
# -*- coding: utf-8 -*-
# Copyright (c) 2012 Giorgos Verigakis <verigak@gmail.com>
#
# 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 = ("○", "◔", "◑", "◕", "●")

42
classification_ScanObjectNN/utils/progress/progress/helpers.py
View file

@ -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):

16
classification_ScanObjectNN/utils/progress/progress/spinner.py
View file

@ -1,5 +1,3 @@
# -*- coding: utf-8 -*-
# Copyright (c) 2012 Giorgos Verigakis <verigak@gmail.com>
#
# 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 = ["⣾", "⣷", "⣯", "⣟", "⡿", "⢿", "⣻", "⣽"]

40
classification_ScanObjectNN/utils/progress/setup.py
View file

@ -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",
],
)

26
classification_ScanObjectNN/utils/progress/test_progress.py
View file

@ -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()

381
part_segmentation/main.py
View file

@ -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)

1
part_segmentation/model/__init__.py
View file

@ -1,2 +1 @@
from __future__ import absolute_import
from .pointMLP import pointMLP

336
part_segmentation/model/pointMLP.py
View file

@ -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]

95
part_segmentation/util/data_util.py
View file

@ -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)

35
part_segmentation/util/util.py
View file

@ -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]

68
pointnet2_ops_lib/pointnet2_ops/pointnet2_modules.py
View file

@ -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

100
pointnet2_ops_lib/pointnet2_ops/pointnet2_utils.py
View file

@ -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

4
pointnet2_ops_lib/setup.py
View file

@ -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,