| import os |
| from tqdm import tqdm |
| import pickle |
| import argparse |
| import pathlib |
| import json |
| import time |
| import torch |
| import torch.nn as nn |
| import torch.nn.parallel |
| import torch.utils.data |
| import numpy as np |
| import torch.nn.functional as F |
| from torch.utils.data import Dataset, DataLoader |
| from metrics import ConfusionMatrix |
| import data_transforms |
| import argparse |
| import random |
| import traceback |
|
|
| """ |
| Model |
| """ |
| class STN3d(nn.Module): |
| def __init__(self, in_channels): |
| super(STN3d, self).__init__() |
| self.conv_layers = nn.Sequential( |
| nn.Conv1d(in_channels, 64, 1), |
| nn.BatchNorm1d(64), |
| nn.ReLU(inplace=True), |
| nn.Conv1d(64, 128, 1), |
| nn.BatchNorm1d(128), |
| nn.ReLU(inplace=True), |
| nn.Conv1d(128, 1024, 1), |
| nn.BatchNorm1d(1024), |
| nn.ReLU(inplace=True) |
| ) |
| self.linear_layers = nn.Sequential( |
| nn.Linear(1024, 512), |
| nn.BatchNorm1d(512), |
| nn.ReLU(inplace=True), |
| nn.Linear(512, 256), |
| nn.BatchNorm1d(256), |
| nn.ReLU(inplace=True), |
| nn.Linear(256, 9) |
| ) |
| self.iden = torch.from_numpy(np.array([1, 0, 0, 0, 1, 0, 0, 0, 1]).astype(np.float32)).reshape(1, 9) |
|
|
| def forward(self, x): |
| batchsize = x.size()[0] |
| x = self.conv_layers(x) |
| x = torch.max(x, 2, keepdim=True)[0] |
| x = x.view(-1, 1024) |
|
|
| x = self.linear_layers(x) |
| iden = self.iden.repeat(batchsize, 1).to(x.device) |
| x = x + iden |
| x = x.view(-1, 3, 3) |
| return x |
|
|
|
|
| class STNkd(nn.Module): |
| def __init__(self, k=64): |
| super(STNkd, self).__init__() |
| self.conv_layers = nn.Sequential( |
| nn.Conv1d(k, 64, 1), |
| nn.BatchNorm1d(64), |
| nn.ReLU(inplace=True), |
| nn.Conv1d(64, 128, 1), |
| nn.BatchNorm1d(128), |
| nn.ReLU(inplace=True), |
| nn.Conv1d(128, 1024, 1), |
| nn.BatchNorm1d(1024), |
| nn.ReLU(inplace=True) |
| ) |
| self.linear_layers = nn.Sequential( |
| nn.Linear(1024, 512), |
| nn.BatchNorm1d(512), |
| nn.ReLU(inplace=True), |
| nn.Linear(512, 256), |
| nn.BatchNorm1d(256), |
| nn.ReLU(inplace=True), |
| nn.Linear(256, k * k) |
| ) |
| self.k = k |
| self.iden = torch.from_numpy(np.eye(self.k).flatten().astype(np.float32)).reshape(1, self.k * self.k) |
|
|
| def forward(self, x): |
| batchsize = x.size()[0] |
| x = self.conv_layers(x) |
| x = torch.max(x, 2, keepdim=True)[0] |
| x = x.view(-1, 1024) |
| x = self.linear_layers(x) |
| iden = self.iden.repeat(batchsize, 1).to(x.device) |
| x = x + iden |
| x = x.view(-1, self.k, self.k) |
| return x |
|
|
|
|
| class PointNetEncoder(nn.Module): |
| def __init__(self, global_feat=True, feature_transform=False, in_channels=3): |
| super(PointNetEncoder, self).__init__() |
| self.stn = STN3d(in_channels) |
| self.conv_layer1 = nn.Sequential( |
| nn.Conv1d(in_channels, 64, 1), |
| nn.BatchNorm1d(64), |
| nn.ReLU(inplace=True), |
| nn.Conv1d(64, 64, 1), |
| nn.BatchNorm1d(64), |
| nn.ReLU(inplace=True) |
| ) |
| self.conv_layer2 = nn.Sequential( |
| nn.Conv1d(64, 64, 1), |
| nn.BatchNorm1d(64), |
| nn.ReLU(inplace=True) |
| ) |
| self.conv_layer3 = nn.Sequential( |
| nn.Conv1d(64, 128, 1), |
| nn.BatchNorm1d(128), |
| nn.ReLU(inplace=True) |
| ) |
| self.conv_layer4 = nn.Sequential( |
| nn.Conv1d(128, 1024, 1), |
| nn.BatchNorm1d(1024) |
| ) |
| self.global_feat = global_feat |
| self.feature_transform = feature_transform |
| if self.feature_transform: |
| self.fstn = STNkd(k=64) |
|
|
| def forward(self, x): |
| B, D, N = x.size() |
| trans = self.stn(x) |
| x = x.transpose(2, 1) |
| if D > 3: |
| feature = x[:, :, 3:] |
| x = x[:, :, :3] |
| x = torch.bmm(x, trans) |
| if D > 3: |
| x = torch.cat([x, feature], dim=2) |
| x = x.transpose(2, 1) |
| x = self.conv_layer1(x) |
|
|
| if self.feature_transform: |
| trans_feat = self.fstn(x) |
| x = x.transpose(2, 1) |
| x = torch.bmm(x, trans_feat) |
| x = x.transpose(2, 1) |
| else: |
| trans_feat = None |
|
|
| pointfeat = x |
| x = self.conv_layer2(x) |
| x = self.conv_layer3(x) |
| x = self.conv_layer4(x) |
| x = torch.max(x, 2, keepdim=True)[0] |
| x = x.view(-1, 1024) |
| |
| |
| graph = construct_graph(x, args.k) |
| context_features = compute_context_aware_features(x, graph) |
| x = x + context_features |
| |
| if self.global_feat: |
| return x, trans, trans_feat |
| else: |
| x = x.view(-1, 1024, 1).repeat(1, 1, N) |
| return torch.cat([x, pointfeat], 1), trans, trans_feat |
|
|
|
|
|
|
| def construct_graph(points, k): |
| """ |
| Construct a dynamic graph where nodes represent points and edges capture semantic similarities. |
| """ |
| |
| dist = torch.cdist(points, points) |
| |
| _, indices = torch.topk(dist, k, largest=False, dim=1) |
| return indices |
|
|
| def compute_context_aware_features(points, graph, normalization_method='mean'): |
| """ |
| Compute context-aware feature adjustments using the constructed graph. |
| """ |
| |
| context_features = torch.zeros_like(points) |
| for i in range(points.size(0)): |
| neighbors = graph[i] |
| if normalization_method == 'mean': |
| context_features[i] = points[neighbors].mean(dim=0) |
| elif normalization_method == 'max': |
| context_features[i] = points[neighbors].max(dim=0)[0] |
| elif normalization_method == 'min': |
| context_features[i] = points[neighbors].min(dim=0)[0] |
| elif normalization_method == 'std': |
| context_features[i] = points[neighbors].std(dim=0) |
| else: |
| raise ValueError("Unknown normalization method: {}".format(normalization_method)) |
| return context_features |
|
|
| def feature_transform_reguliarzer(trans): |
| d = trans.size()[1] |
| I = torch.eye(d)[None, :, :] |
| if trans.is_cuda: |
| I = I.cuda() |
| loss = torch.mean(torch.norm(torch.bmm(trans, trans.transpose(2, 1)) - I, dim=(1, 2))) |
| return loss |
|
|
| class Model(nn.Module): |
| def __init__(self, in_channels=3, num_classes=40, scale=0.001): |
| super().__init__() |
| self.mat_diff_loss_scale = scale |
| self.backbone = PointNetEncoder(global_feat=True, feature_transform=True, in_channels=in_channels) |
| self.cls_head = nn.Sequential( |
| nn.Linear(1024, 512), |
| nn.BatchNorm1d(512), |
| nn.ReLU(inplace=True), |
| nn.Linear(512, 256), |
| nn.Dropout(p=0.4), |
| nn.BatchNorm1d(256), |
| nn.ReLU(inplace=True), |
| nn.Linear(256, num_classes) |
| ) |
| |
| def forward(self, x, gts): |
| x, trans, trans_feat = self.backbone(x) |
| x = self.cls_head(x) |
| x = F.log_softmax(x, dim=1) |
| loss = F.nll_loss(x, gts) |
| mat_diff_loss = feature_transform_reguliarzer(trans_feat) |
| total_loss = loss + mat_diff_loss * self.mat_diff_loss_scale |
| return total_loss, x |
|
|
|
|
| """ |
| dataset and normalization |
| """ |
| 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 |
|
|
|
|
| class ModelNetDataset(Dataset): |
| def __init__(self, data_root, num_category, num_points, split='train'): |
| self.root = data_root |
| self.npoints = num_points |
| self.uniform = True |
| self.use_normals = True |
| self.num_category = num_category |
|
|
| if self.num_category == 10: |
| self.catfile = os.path.join(self.root, 'modelnet10_shape_names.txt') |
| else: |
| self.catfile = os.path.join(self.root, 'modelnet40_shape_names.txt') |
|
|
| self.cat = [line.rstrip() for line in open(self.catfile)] |
| self.classes = dict(zip(self.cat, range(len(self.cat)))) |
|
|
| shape_ids = {} |
| if self.num_category == 10: |
| shape_ids['train'] = [line.rstrip() for line in open(os.path.join(self.root, 'modelnet10_train.txt'))] |
| shape_ids['test'] = [line.rstrip() for line in open(os.path.join(self.root, 'modelnet10_test.txt'))] |
| else: |
| shape_ids['train'] = [line.rstrip() for line in open(os.path.join(self.root, 'modelnet40_train.txt'))] |
| shape_ids['test'] = [line.rstrip() for line in open(os.path.join(self.root, 'modelnet40_test.txt'))] |
|
|
| assert (split == 'train' or split == 'test') |
| shape_names = ['_'.join(x.split('_')[0:-1]) for x in shape_ids[split]] |
| self.datapath = [(shape_names[i], os.path.join(self.root, shape_names[i], shape_ids[split][i]) + '.txt') for i |
| in range(len(shape_ids[split]))] |
| print('The size of %s data is %d' % (split, len(self.datapath))) |
|
|
| if self.uniform: |
| self.data_path = os.path.join(data_root, 'modelnet%d_%s_%dpts_fps.dat' % (self.num_category, split, self.npoints)) |
| else: |
| self.data_path = os.path.join(data_root, 'modelnet%d_%s_%dpts.dat' % (self.num_category, split, self.npoints)) |
|
|
| print('Load processed data from %s...' % self.data_path) |
| with open(self.data_path, 'rb') as f: |
| self.list_of_points, self.list_of_labels = pickle.load(f) |
|
|
| def __len__(self): |
| return len(self.datapath) |
|
|
| def __getitem__(self, index): |
| point_set, label = self.list_of_points[index], self.list_of_labels[index] |
| point_set[:, 0:3] = pc_normalize(point_set[:, 0:3]) |
| if not self.use_normals: |
| point_set = point_set[:, 0:3] |
| return point_set, label[0] |
|
|
|
|
| def seed_everything(seed=11): |
| random.seed(seed) |
| np.random.seed(seed) |
| torch.manual_seed(seed) |
| torch.cuda.manual_seed(seed) |
| torch.cuda.manual_seed_all(seed) |
| torch.backends.cudnn.deterministic = True |
| torch.backends.cudnn.benchmark = False |
|
|
|
|
| def main(args): |
|
|
| seed_everything(args.seed) |
|
|
| final_infos = {} |
| all_results = {} |
|
|
| pathlib.Path(args.out_dir).mkdir(parents=True, exist_ok=True) |
| |
| datasets, dataloaders = {}, {} |
| for split in ['train', 'test']: |
| datasets[split] = ModelNetDataset(args.data_root, args.num_category, args.num_points, split) |
| dataloaders[split] = DataLoader(datasets[split], batch_size=args.batch_size, shuffle=(split == 'train'), |
| drop_last=(split == 'train'), num_workers=8) |
| |
| model = Model(in_channels=args.in_channels).cuda() |
| optimizer = torch.optim.Adam( |
| model.parameters(), lr=args.learning_rate, |
| betas=(0.9, 0.999), eps=1e-8, |
| weight_decay=1e-4 |
| ) |
| scheduler = torch.optim.lr_scheduler.StepLR( |
| optimizer, step_size=20, gamma=0.7 |
| ) |
| train_losses = [] |
| print("Training model...") |
| model.train() |
| global_step = 0 |
| cur_epoch = 0 |
| best_oa = 0 |
| best_acc = 0 |
|
|
| start_time = time.time() |
| for epoch in tqdm(range(args.max_epoch), desc='training'): |
| model.train() |
| cm = ConfusionMatrix(num_classes=len(datasets['train'].classes)) |
| for points, target in tqdm(dataloaders['train'], desc=f'epoch {cur_epoch}/{args.max_epoch}'): |
| |
| points = points.data.numpy() |
| points = data_transforms.random_point_dropout(points) |
| points[:, :, 0:3] = data_transforms.random_scale_point_cloud(points[:, :, 0:3]) |
| points[:, :, 0:3] = data_transforms.shift_point_cloud(points[:, :, 0:3]) |
| points = torch.from_numpy(points).transpose(2, 1).contiguous() |
| |
| points, target = points.cuda(), target.long().cuda() |
| |
| loss, logits = model(points, target) |
| loss.backward() |
|
|
| torch.nn.utils.clip_grad_norm_(model.parameters(), 1, norm_type=2) |
| optimizer.step() |
| model.zero_grad() |
| |
| |
| logs = {"loss": loss.detach().item()} |
| train_losses.append(loss.detach().item()) |
| cm.update(logits.argmax(dim=1), target) |
| |
| scheduler.step() |
| end_time = time.time() |
| training_time = end_time - start_time |
| macc, overallacc, accs = cm.all_acc() |
| print(f"iter: {global_step}/{args.max_epoch*len(dataloaders['train'])}, \ |
| train_macc: {macc}, train_oa: {overallacc}") |
| |
| if (cur_epoch % args.val_per_epoch == 0 and cur_epoch != 0) or cur_epoch == (args.max_epoch - 1): |
| model.eval() |
| cm = ConfusionMatrix(num_classes=datasets['test'].num_category) |
| pbar = tqdm(enumerate(dataloaders['test']), total=dataloaders['test'].__len__()) |
| |
| for idx, (points, target) in pbar: |
| points, target = points.cuda(), target.long().cuda() |
| points = points.transpose(2, 1).contiguous() |
| loss, logits = model(points, target) |
| cm.update(logits.argmax(dim=1), target) |
| |
| tp, count = cm.tp, cm.count |
| macc, overallacc, accs = cm.cal_acc(tp, count) |
| print(f"iter: {global_step}/{args.max_epoch*len(dataloaders['train'])}, \ |
| val_macc: {macc}, val_oa: {overallacc}") |
| |
| if overallacc > best_oa: |
| best_oa = overallacc |
| best_acc = macc |
| best_epoch = cur_epoch |
| torch.save(model.state_dict(), os.path.join(args.out_dir, 'best.pth')) |
| cur_epoch += 1 |
|
|
| print(f"finish epoch {cur_epoch} training") |
|
|
| final_infos = { |
| "modelnet" + str(args.num_category):{ |
| "means":{ |
| "best_oa": best_oa, |
| "best_acc": best_acc, |
| "epoch": best_epoch |
| } |
| } |
| } |
| with open(os.path.join(args.out_dir, "final_info.json"), "w") as f: |
| json.dump(final_infos, f) |
|
|
| if __name__ == "__main__": |
|
|
| parser = argparse.ArgumentParser() |
| parser.add_argument("--batch_size", type=int, default=64) |
| parser.add_argument("--out_dir", type=str, default="run_0") |
| parser.add_argument("--in_channels", type=int, default=6) |
| parser.add_argument("--num_points", type=int, default=1024) |
| parser.add_argument("--num_category", type=int, choices=[10, 40], default=40) |
| parser.add_argument("--data_root", type=str, default='./datasets/modelnet40') |
| parser.add_argument("--learning_rate", type=float, default=1e-3) |
| parser.add_argument("--max_epoch", type=int, default=200) |
| parser.add_argument("--val_per_epoch", type=int, default=5) |
| parser.add_argument("--k", type=int, default=5, help="Number of neighbors for graph construction") |
| parser.add_argument("--seed", type=int, default=666) |
| args = parser.parse_args() |
|
|
| try: |
| main(args) |
| except Exception as e: |
| print("Original error in subprocess:", flush=True) |
| traceback.print_exc(file=open(os.path.join(args.out_dir, "traceback.log"), "w")) |
| raise |
