def transform_setup(graph_u=False, graph_gcn=False, rotation=180, samplePoints=1024, mesh=False, node_translation=0.01): if not graph_u and not graph_gcn: # Default transformation for scale noralization, centering, point sampling and rotating pretransform = T.Compose([T.NormalizeScale(), T.Center()]) transform = T.Compose([ T.SamplePoints(samplePoints), T.RandomRotate(rotation[0], rotation[1]) ]) print("pointnet rotation {}".format(rotation)) elif graph_u: pretransform = T.Compose([T.NormalizeScale(), T.Center()]) transform = T.Compose([ T.NormalizeScale(), T.Center(), T.SamplePoints(samplePoints, True, True), T.RandomRotate(rotation[0], rotation[1]), T.KNNGraph(k=graph_u) ]) elif graph_gcn: pretransform = T.Compose([T.NormalizeScale(), T.Center()]) if mesh: if mesh == "extraFeatures": transform = T.Compose([ T.RandomRotate(rotation[0], rotation[1]), T.GenerateMeshNormals(), T.FaceToEdge(True), T.Distance(norm=True), T.TargetIndegree(cat=True) ]) # , else: transform = T.Compose([ T.RandomRotate(rotation[0], rotation[1]), T.GenerateMeshNormals(), T.FaceToEdge(True), T.Distance(norm=True), T.TargetIndegree(cat=True) ]) else: transform = T.Compose([ T.SamplePoints(samplePoints, True, True), T.KNNGraph(k=graph_gcn), T.Distance(norm=True) ]) print("no mesh") print("Rotation {}".format(rotation)) print("Meshing {}".format(mesh)) else: print('no transfom') return transform, pretransform
def main(): opt = OptInit().get_args() logging.info('===> Creating dataloader ...') train_dataset = GeoData.S3DIS(opt.data_dir, opt.area, True, pre_transform=T.NormalizeScale()) train_loader = DenseDataLoader(train_dataset, batch_size=opt.batch_size, shuffle=True, num_workers=4) test_dataset = GeoData.S3DIS(opt.data_dir, opt.area, train=False, pre_transform=T.NormalizeScale()) test_loader = DenseDataLoader(test_dataset, batch_size=opt.batch_size, shuffle=False, num_workers=0) opt.n_classes = train_loader.dataset.num_classes logging.info('===> Loading the network ...') model = DenseDeepGCN(opt).to(opt.device) if opt.multi_gpus: model = DataParallel(DenseDeepGCN(opt)).to(opt.device) logging.info('===> loading pre-trained ...') model, opt.best_value, opt.epoch = load_pretrained_models( model, opt.pretrained_model, opt.phase) logging.info(model) logging.info('===> Init the optimizer ...') criterion = torch.nn.CrossEntropyLoss().to(opt.device) optimizer = torch.optim.Adam(model.parameters(), lr=opt.lr) scheduler = torch.optim.lr_scheduler.StepLR(optimizer, opt.lr_adjust_freq, opt.lr_decay_rate) optimizer, scheduler, opt.lr = load_pretrained_optimizer( opt.pretrained_model, optimizer, scheduler, opt.lr) logging.info('===> Init Metric ...') opt.losses = AverageMeter() opt.test_value = 0. logging.info('===> start training ...') for _ in range(opt.epoch, opt.total_epochs): opt.epoch += 1 logging.info('Epoch:{}'.format(opt.epoch)) train(model, train_loader, optimizer, scheduler, criterion, opt) if opt.epoch % opt.eval_freq == 0 and opt.eval_freq != -1: test(model, test_loader, opt) scheduler.step() logging.info('Saving the final model.Finish!')
def __init__(self, train): dataset = "ModelNet40" pre_transform, transform = T.NormalizeScale(), T.SamplePoints(1024) path = osp.join(osp.dirname(osp.realpath(__file__)), "../..", "data", dataset) if not osp.exists(path): ModelNet(path, "40", transform, pre_transform) super(ModelNet40, self).__init__(path, name="40", train=train, transform=transform, pre_transform=pre_transform)
def load_modelnet(version='10', point_flag=False): """ :param point_flag: Sample points if point_flag true. Otherwise load mesh :return: train_dataset, test_dataset """ assert version in ['10', '40'] if point_flag: pre_transform, transform = T.NormalizeScale(), T.SamplePoints(1024) else: pre_transform, transform = FaceToEdge(), None # path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data/ModelNet' + version) path = '/home/cai.507/Documents/DeepLearning/ModelNet' + version train_dataset = ModelNet(path, version, True, transform=transform, pre_transform=pre_transform) test_dataset = ModelNet(path, version, False, transform=transform, pre_transform=pre_transform) return train_dataset, test_dataset
class MNISTSuperpixels(LightningDataModule): def __init__( self, data_dir: str = "data/", batch_size: int = 32, num_workers: int = 0, pin_memory: bool = False, train_val_test_split: Sequence[int] = (55_000, 5_000, 10_000), n_segments: int = 75, k: int = 10, loop: bool = True, **kwargs, ): super().__init__() self.data_dir = data_dir self.batch_size = batch_size self.num_workers = num_workers self.pin_memory = pin_memory self.train_val_test_split = train_val_test_split self.n_segments = n_segments self.k = k self.loop = loop self.slic_kwargs = kwargs assert 1 <= n_segments <= 28 * 28 self.pre_transform = T.Compose([ T.NormalizeScale(), ]) self.transform = None self.data_train: Optional[Dataset] = None self.data_val: Optional[Dataset] = None self.data_test: Optional[Dataset] = None
def __init__(self, dataset_opt, training_opt): super().__init__(dataset_opt, training_opt) self._data_path = os.path.join(dataset_opt.dataroot, "ShapeNet") try: self._category = dataset_opt.category except KeyError: self._category = None pre_transform = T.NormalizeScale() train_transform = T.Compose( [T.FixedPoints(dataset_opt.num_points), RandomNoise()]) test_transform = T.FixedPoints(dataset_opt.num_points) train_dataset = ShapeNet( self._data_path, self._category, include_normals=dataset_opt.normal, split="trainval", pre_transform=pre_transform, transform=train_transform, ) test_dataset = ShapeNet( self._data_path, self._category, include_normals=dataset_opt.normal, split="test", pre_transform=pre_transform, transform=test_transform, ) self._categories = train_dataset.categories self._create_dataloaders(train_dataset, test_dataset)
def augment_transforms(args): """ define transformation """ pre_transform = None if args.norm == 'scale': pre_transform = T.NormalizeScale() elif args.norm == 'bbox': pre_transform = NormalizeBox() elif args.norm == 'sphere': pre_transform = NormalizeSphere(center=True) elif args.norm == 'sphere_wo_center': pre_transform = NormalizeSphere(center=False) else: pass transform = [] # Shapenet if args.task == 'segmentation': transform.append(T.FixedPoints(args.num_pts)) # Modelnet if args.task == 'classification': transform.append(T.SamplePoints(args.num_pts)) transform = T.Compose(transform) return pre_transform, transform
def __init__(self, root, **kwargs): self.root = os.path.join(root, 'ModelNet40') self.name = '40' self.pre_transform = T.NormalizeScale() self.transform = T.SamplePoints(1024) self.label_parser = lambda data: data.y self.num_classes = 40
def __init__(self, dataset_opt, training_opt): super().__init__(dataset_opt, training_opt) number = dataset_opt.number if str(number) not in AVAILABLE_NUMBERS: raise Exception("Only ModelNet10 and ModelNet40 are available") name = "ModelNet{}".format(number) self._data_path = osp.join(osp.dirname(osp.realpath(__file__)), "..", "data", name) pre_transform = T.Compose([T.NormalizeScale(), MeshToNormal()]) transform = (T.SamplePoints(dataset_opt.num_points) if contains_key( dataset_opt, "num_points") else None) train_dataset = ModelNet( self._data_path, name=str(number), train=True, transform=transform, pre_transform=pre_transform, ) test_dataset = ModelNet( self._data_path, name=str(number), train=False, transform=transform, pre_transform=pre_transform, ) self._create_dataloaders(train_dataset, test_dataset, validation=None)
def get_view_transform(self, k, num_pts): R = rotation_matrix(np.pi / 3., 0., np.pi / 6. * k) transformation = TG.Compose([ TG.NormalizeScale(), TG.LinearTransformation(R), TG.SamplePoints(num=num_pts, include_normals=self.generate_norms) ]) return transformation
def __init__(self, root='data/ShapeNet', train=True, categories=None, include_normals=True, split='trainval', transform=None, pre_transform=None, pre_filter=None, repeat_to=None): # Modified here to add repeat_to if categories is None: categories = list(self.category_ids.keys()) if isinstance(categories, str): categories = [categories] assert all(category in self.category_ids for category in categories) self.categories = categories # Default settings pre_transform = T.NormalizeScale() pre_filter = None include_normals = True if train: transform = T.Compose([ T.RandomTranslate(0.01), T.RandomRotate(15, axis=0), T.RandomRotate(15, axis=1), T.RandomRotate(15, axis=2) ]) split = 'trainval' else: transform = None split = 'test' super().__init__(root, transform, pre_transform, pre_filter) # Modified here to add repeat_to if split == 'train': path = self.processed_paths[0] elif split == 'val': path = self.processed_paths[1] elif split == 'test': path = self.processed_paths[2] elif split == 'trainval': path = self.processed_paths[3] else: raise ValueError((f'Split {split} found, but expected either ' 'train, val, trainval or test')) self.data, self.slices = torch.load(path) self.data.x = self.data.x if include_normals else None self.y_mask = torch.zeros((len(self.seg_classes.keys()), 50), dtype=torch.bool) for i, labels in enumerate(self.seg_classes.values()): self.y_mask[i, labels] = 1 self.repeat_to = repeat_to # Modified here to add repeat_to
def get_s3dis_dataloaders(root_dir, phases, batch_size, category=5, augment=False): """ Create Dataset and Dataloader classes of the S3DIS dataset, for the phases required (`train`, `test`). :param root_dir: Directory with the h5 files :param phases: List of phases. Should be from {`train`, `test`} :param batch_size: Batch size :param category: Area used for test set (1, 2, 3, 4, 5, or 6) :return: 2 dictionaries, each containing Dataset or Dataloader for all phases """ datasets = { 'train': S3DIS(root_dir, category, True, pre_transform=T.NormalizeScale()), 'test': S3DIS(root_dir, category, False, pre_transform=T.NormalizeScale()) } dataloaders = {x: DenseDataLoader(datasets[x], batch_size=batch_size, num_workers=4, shuffle=(x == 'train')) for x in phases} return datasets, dataloaders, datasets['train'].num_classes
def main(): opt = OptInit().initialize() opt.printer.info('===> Creating dataloader ...') train_dataset = GeoData.S3DIS(opt.train_path, 5, True, pre_transform=T.NormalizeScale()) if opt.multi_gpus: train_loader = DataListLoader(train_dataset, batch_size=opt.batch_size, shuffle=True, num_workers=4) else: train_loader = DataLoader(train_dataset, batch_size=opt.batch_size, shuffle=True, num_workers=4) opt.n_classes = train_loader.dataset.num_classes opt.printer.info('===> Loading the network ...') model = SparseDeepGCN(opt).to(opt.device) if opt.multi_gpus: model = DataParallel(SparseDeepGCN(opt)).to(opt.device) opt.printer.info('===> loading pre-trained ...') model, opt.best_value, opt.epoch = load_pretrained_models( model, opt.pretrained_model, opt.phase) opt.printer.info('===> Init the optimizer ...') criterion = torch.nn.CrossEntropyLoss().to(opt.device) if opt.optim.lower() == 'adam': optimizer = torch.optim.Adam(model.parameters(), lr=opt.lr) elif opt.optim.lower() == 'radam': optimizer = optim.RAdam(model.parameters(), lr=opt.lr) else: raise NotImplementedError('opt.optim is not supported') scheduler = torch.optim.lr_scheduler.StepLR(optimizer, opt.lr_adjust_freq, opt.lr_decay_rate) optimizer, scheduler, opt.lr = load_pretrained_optimizer( opt.pretrained_model, optimizer, scheduler, opt.lr) opt.printer.info('===> Init Metric ...') opt.losses = AverageMeter() # opt.test_metric = miou # opt.test_values = AverageMeter() opt.test_value = 0. opt.printer.info('===> start training ...') for _ in range(opt.total_epochs): opt.epoch += 1 train(model, train_loader, optimizer, scheduler, criterion, opt) # test_value = test(model, test_loader, test_metric, opt) scheduler.step() opt.printer.info('Saving the final model.Finish!')
def __init__(self, root='data/ModelNet', name='40', train=True): # Default setting pre_transform = T.NormalizeScale() transform = T.SamplePoints(1024) pre_filter = None super().__init__(root + name, name=name, train=train, transform=transform, pre_transform=pre_transform, pre_filter=pre_filter)
def get_dataset(num_points): name = 'ModelNet10' path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', name) pre_transform = T.NormalizeScale() transform = T.SamplePoints(num_points) train_dataset = ModelNet(path, name='10', train=True, transform=transform, pre_transform=pre_transform) test_dataset = ModelNet(path, name='10', train=False, transform=transform, pre_transform=pre_transform) return train_dataset, test_dataset
def __init__(self, root, category, **kwargs): self.root = os.path.join(root, 'ShapeNet') self.category = category self.categories = [category] self.pre_transform = T.NormalizeScale() self.train_transform = T.Compose([ T.RandomTranslate(0.01), T.RandomRotate(15, axis=0), T.RandomRotate(15, axis=1), T.RandomRotate(15, axis=2), ]) self.label_parser = lambda data: data.y
def predict_age(file_path="/media/original/data/vtps/sub-CC00050XX01_ses-7201_hemi-L_inflated_reduce50.vtp"): torch.manual_seed(0) if osp.isfile(file_path): # mesh = read(file_path) # reader = vtk.vtkPolyDataReader() reader = vtk.vtkXMLPolyDataReader() reader.SetFileName(file_path) reader.Update() # output = reader.GetOutput() points = torch.tensor(np.array(reader.GetOutput().GetPoints().GetData())) local_features = ['corrected_thickness', 'curvature', 'sulcal_depth'] x = get_features(local_features, reader) transform = T.NormalizeScale() # transform_samp = T.FixedPoints(10000) data = Data(batch=torch.zeros_like(x[:, 0]).long(), x=x, pos=points) data = transform(data) # data = transform_samp(data) # data = Data(batch=torch.zeros_like(x[:, 0]).long(), x=x, pos=points) # data = Data(x=x, pos=points) try: nvidia_smi.nvmlInit() handle = nvidia_smi.nvmlDeviceGetHandleByIndex(0) mem_res = nvidia_smi.nvmlDeviceGetMemoryInfo(handle) free_mem = mem_res.free / 1024 ** 2 except: free_mem = 0 device = torch.device('cuda' if torch.cuda.is_available() and free_mem >= GPU_MEM_LIMIT else 'cpu') numb_local_features = x.size(1) numb_global_features = 0 model = Net(numb_local_features, numb_global_features).to(device) model.load_state_dict(torch.load(MODEL_PATH, map_location=device)) model.eval() # data_loader = DataLoader([data], batch_size=1, shuffle=False) # print(len(data_loader)) # pred = model(next(iter(data_loader)).to(device)) pred = model(data.to(device)) return pred.item() else: return 'Unable to predict..'
def main(): opt = OptInit().initialize() opt.batch_size = 1 os.environ["CUDA_VISIBLE_DEVICES"] = opt.gpuNum print('===> Creating dataloader...') # def __init__(self, # root, # is_train=True, # is_validation=False, # is_test=False, # num_channel=5, # pre_transform=None, # pre_filter=None) test_dataset = BigredDataset(root=opt.test_path, is_train=False, is_validation=False, is_test=True, num_channel=5, pre_transform=T.NormalizeScale()) test_loader = DenseDataLoader(test_dataset, batch_size=opt.batch_size, shuffle=False, num_workers=32) opt.n_classes = 2 print('len(test_loader):', len(test_loader)) print('phase: ', opt.phase) print('batch_size: ', opt.batch_size) print('use_cpu: ', opt.use_cpu) print('gpuNum: ', opt.gpuNum) print('multi_gpus: ', opt.multi_gpus) print('test_path: ', opt.test_path) print('in_channels: ', opt.in_channels) print('device: ', opt.device) print('===> Loading the network ...') model = DenseDeepGCN(opt).to(opt.device) load_package = torch.load(opt.pretrained_model) model, opt.best_value, opt.epoch = load_pretrained_models( model, opt.pretrained_model, opt.phase) pdb.set_trace() for item in load_package.keys(): if (item != 'optimizer_state_dict' and item != 'state_dict' and item != 'scheduler_state_dict'): print(str(item), load_package[item]) print('===> Start Evaluation ...') test(model, test_loader, opt)
def __init__(self, root, classification=False, **kwargs): self.root = os.path.join(root, 'ShapeNet') self.pre_transform = T.NormalizeScale() self.train_transform = T.Compose([ RandomSamplePoints(2048), T.RandomTranslate(0.01), T.RandomRotate(15, axis=0), T.RandomRotate(15, axis=1), T.RandomRotate(15, axis=2), ]) self.val_transform = RandomSamplePoints(2048) self.num_classes = len(self.categories) if not classification: self.label_parser = lambda data: data.y else: self.label_parser = lambda data: data.cid
def main(): opt = OptInit().initialize() print('===> Creating dataloader...') test_dataset = GeoData.S3DIS(opt.test_path, 5, False, pre_transform=T.NormalizeScale()) test_loader = DataLoader(test_dataset, batch_size=opt.batch_size, shuffle=False, num_workers=0) opt.n_classes = test_loader.dataset.num_classes if opt.no_clutter: opt.n_classes -= 1 print('===> Loading the network ...') model = getattr(models, opt.model_name)(opt).to(opt.device) model, opt.best_value, opt.epoch = load_pretrained_models(model, opt.pretrained_model, opt.phase) print('===> Start Evaluation ...') test(opt.model, test_loader, opt)
def __init__(self, dataset_opt, training_opt): super().__init__(dataset_opt, training_opt) self._data_path = os.path.join(dataset_opt.dataroot, 'ShapeNet') self._category = dataset_opt.shapenet.category transform = T.Compose([ T.RandomTranslate(0.01), T.RandomRotate(15, axis=0), T.RandomRotate(15, axis=1), T.RandomRotate(15, axis=2) ]) pre_transform = T.NormalizeScale() train_dataset = ShapeNet(self._data_path, self._category, train=True, transform=transform, pre_transform=pre_transform) test_dataset = ShapeNet( self._data_path, self._category, train=False, pre_transform=pre_transform) self.create_dataloaders(train_dataset, test_dataset, validation=None)
def main(): opt = OptInit().initialize() print('===> Creating dataloader ...') train_dataset = GeoData.S3DIS(opt.train_path, 5, True, pre_transform=T.NormalizeScale()) if opt.multi_gpus: train_loader = DataListLoader(train_dataset, batch_size=opt.batch_size, shuffle=True, num_workers=4) else: train_loader = DataLoader(train_dataset, batch_size=opt.batch_size, shuffle=True, num_workers=4) opt.n_classes = train_loader.dataset.num_classes print('===> Loading the network ...') opt.model = getattr(models, opt.model_name)(opt).to(opt.device) if opt.multi_gpus: opt.model = DataParallel(getattr(models, opt.model_name)(opt)).to(opt.device) print('===> loading pre-trained ...') load_pretrained_models(opt) print('===> Init the optimizer ...') opt.criterion = torch.nn.CrossEntropyLoss().to(opt.device) opt.valid_metric = miou opt.optimizer = torch.optim.Adam(opt.model.parameters(), lr=opt.lr) opt.scheduler = torch.optim.lr_scheduler.StepLR(opt.optimizer, opt.lr_adjust_freq, 0.5) load_pretrained_optimizer(opt) print('===> start training ...') for _ in range(opt.total_epochs): opt.epoch += 1 train(train_loader, opt) # valid(train_loader, opt) opt.scheduler.step() print('Saving the final model.Finish!')
def get_dataloader(num_points, b_size, name='10'): path = 'ModelNet' + name pre_transform = T.NormalizeScale() transform = T.SamplePoints(num_points) train_dataset = ModelNet('dataset/' + path, name=name, train=True, transform=transform, pre_transform=pre_transform) test_dataset = ModelNet('dataset/' + path, name=name, train=False, transform=transform, pre_transform=pre_transform) train_loader = DataLoader(train_dataset, batch_size=b_size, shuffle=True) test_loader = DataLoader(test_dataset, batch_size=b_size, shuffle=False) return train_loader, test_loader
def main(epochs, lr, classification, pool, heatmap): print('pool', pool) data_path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', 'coseg', 'vases') arguments = 'c{}_e{}_lr{}_p{}'.format(classification, epochs, lr, '-'.join([str(x) for x in pool])) run_path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'runs', arguments) if os.path.exists(run_path): shutil.rmtree(run_path) os.makedirs(run_path) # pre_filter = lambda x: x.shape_id.item() < 10 pre_transform = T.Compose( [T.NormalizeScale(), FaceToEdge(remove_faces=False), CalcEdgeFeaturesTransform(), AddMeshStructureTransform()]) train_dataset = COSEG(data_path, classification=classification, train=True, pre_transform=pre_transform) test_dataset = COSEG(data_path, classification=classification, train=False, pre_transform=pre_transform) train_loader = DataLoader(train_dataset, batch_size=1, shuffle=True) test_loader = DataLoader(test_dataset, batch_size=1) device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') print(device) model = Net(train_dataset.num_classes, pool, run_path, heatmap).to(device) train_writer = SummaryWriter(run_path + '/train') test_writer = SummaryWriter(run_path + '/test') optimizer = torch.optim.Adam(model.parameters(), lr=lr) visualize_gt(10, train_loader, test_loader, run_path, heatmap) for epoch in range(epochs): train_acc, train_loss = train(model, device, optimizer, train_loader) train_writer.add_scalar('accuracy', train_acc, epoch) train_writer.add_scalar('loss', train_loss, epoch) test_acc, test_loss = test(model, device, test_loader) test_writer.add_scalar('accuracy', test_acc, epoch) test_writer.add_scalar('loss', test_loss, epoch) print('epoch: {:04d}, train acc: {:.4f}, test acc: {:.4f}'.format(epoch + 1, train_acc, test_acc)) train_writer.close() test_writer.close() visualize_pred(10, train_loader, test_loader, run_path, model, device)
def main(): opt = OptInit().get_args() logging.info('===> Creating dataloader...') test_dataset = GeoData.S3DIS(opt.data_dir, opt.area, train=False, pre_transform=T.NormalizeScale()) test_loader = DenseDataLoader(test_dataset, batch_size=opt.batch_size, shuffle=False, num_workers=0) opt.n_classes = test_loader.dataset.num_classes if opt.no_clutter: opt.n_classes -= 1 logging.info('===> Loading the network ...') model = DenseDeepGCN(opt).to(opt.device) model, opt.best_value, opt.epoch = load_pretrained_models( model, opt.pretrained_model, opt.phase) logging.info('===> Start Evaluation ...') test(model, test_loader, opt)
def __init__(self, data_dir, batch_size, shuffle, validation_split, num_workers, num_points, training=True): # trsfm = transforms.Compose([ # transforms.ToTensor(), # transforms.Normalize((0.1307,), (0.3081,)) # ]) self.data_dir = data_dir path = osp.join(self.data_dir, 'ModelNet10') pre_transform, transform = T.NormalizeScale(), T.SamplePoints( num_points) train_dataset = ModelNet(path, '10', training, transform, pre_transform) super(MyModelNetDataLoader, self).__init__(train_dataset, batch_size=batch_size, shuffle=shuffle)
def augment_transforms(args): """ define transformation """ pre_transform = None if args.norm == 'scale': pre_transform = T.NormalizeScale() elif args.norm == 'sphere': pre_transform = NormalizeSphere(center=True) elif args.norm == 'sphere_wo_center': pre_transform = NormalizeSphere(center=False) else: pass transform = [] if args.dataset == 'shapenet': transform.append(T.FixedPoints(args.num_pts)) if args.dataset == 'modelnet': transform.append(T.SamplePoints(args.num_pts)) # if args.is_randRotY: # transform.append(T.RandomRotate(180, axis=1)) transform = T.Compose(transform) return pre_transform, transform
from torchmetrics.functional import jaccard_index import torch_geometric.transforms as T from torch_geometric.datasets import ShapeNet from torch_geometric.loader import DataLoader from torch_geometric.nn import MLP, DynamicEdgeConv category = 'Airplane' # Pass in `None` to train on all categories. path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', 'ShapeNet') transform = T.Compose([ T.RandomTranslate(0.01), T.RandomRotate(15, axis=0), T.RandomRotate(15, axis=1), T.RandomRotate(15, axis=2) ]) pre_transform = T.NormalizeScale() train_dataset = ShapeNet(path, category, split='trainval', transform=transform, pre_transform=pre_transform) test_dataset = ShapeNet(path, category, split='test', pre_transform=pre_transform) train_loader = DataLoader(train_dataset, batch_size=10, shuffle=True, num_workers=6) test_loader = DataLoader(test_dataset, batch_size=10, shuffle=False, num_workers=6) class Net(torch.nn.Module): def __init__(self, out_channels, k=30, aggr='max'): super().__init__() self.conv1 = DynamicEdgeConv(MLP([2 * 6, 64, 64]), k, aggr)
def main(): if not torch.cuda.is_available(): logging.info('no gpu device available') sys.exit(1) if args.random_seed: args.seed = np.random.randint(0, 1000, 1) np.random.seed(args.seed) torch.cuda.set_device(args.gpu) cudnn.benchmark = True torch.manual_seed(args.seed) cudnn.enabled = True torch.cuda.manual_seed(args.seed) logging.info('gpu device = %d' % args.gpu) logging.info("args = %s", args) # dataset modelnet pre_transform, transform = T.NormalizeScale(), T.SamplePoints( args.num_points) train_dataset = GeoData.ModelNet(os.path.join(args.data, 'modelnet10'), '10', True, transform, pre_transform) train_queue = DenseDataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=args.batch_size // 2) test_dataset = GeoData.ModelNet(os.path.join(args.data, 'modelnet10'), '10', False, transform, pre_transform) valid_queue = DenseDataLoader(test_dataset, batch_size=args.batch_size, shuffle=True, num_workers=args.batch_size // 2) n_classes = train_queue.dataset.num_classes criterion = torch.nn.CrossEntropyLoss().cuda() model = Network(args.init_channels, n_classes, args.num_cells, criterion, args.n_steps, in_channels=args.in_channels, emb_dims=args.emb_dims, dropout=args.dropout, k=args.k).cuda() logging.info("param size = %fMB", utils.count_parameters_in_MB(model)) num_edges = model._steps * 2 post_train = 5 # import pdb;pdb.set_trace() args.epochs = args.warmup_dec_epoch + args.decision_freq * ( num_edges - 1) + post_train + 1 logging.info("total epochs: %d", args.epochs) optimizer = torch.optim.SGD(model.parameters(), args.learning_rate, momentum=args.momentum, weight_decay=args.weight_decay) scheduler = torch.optim.lr_scheduler.CosineAnnealingLR( optimizer, float(args.epochs), eta_min=args.learning_rate_min) architect = Architect(model, args) normal_selected_idxs = torch.tensor(len(model.alphas_normal) * [-1], requires_grad=False, dtype=torch.int).cuda() normal_candidate_flags = torch.tensor(len(model.alphas_normal) * [True], requires_grad=False, dtype=torch.bool).cuda() logging.info('normal_selected_idxs: {}'.format(normal_selected_idxs)) logging.info('normal_candidate_flags: {}'.format(normal_candidate_flags)) model.normal_selected_idxs = normal_selected_idxs model.normal_candidate_flags = normal_candidate_flags print(F.softmax(torch.stack(model.alphas_normal, dim=0), dim=-1).detach()) count = 0 normal_probs_history = [] train_losses, valid_losses = utils.AverageMeter(), utils.AverageMeter() for epoch in range(args.epochs): lr = scheduler.get_lr()[0] logging.info('epoch %d lr %e', epoch, lr) # training # import pdb;pdb.set_trace() att = model.show_att() beta = model.show_beta() train_acc, train_losses = train(train_queue, valid_queue, model, architect, criterion, optimizer, lr, train_losses) valid_overall_acc, valid_class_acc, valid_losses = infer( valid_queue, model, criterion, valid_losses) logging.info( 'train_acc %f\tvalid_overall_acc %f \t valid_class_acc %f', train_acc, valid_overall_acc, valid_class_acc) logging.info('beta %s', beta.cpu().detach().numpy()) logging.info('att %s', att.cpu().detach().numpy()) # make edge decisions saved_memory_normal, model.normal_selected_idxs, \ model.normal_candidate_flags = edge_decision('normal', model.alphas_normal, model.normal_selected_idxs, model.normal_candidate_flags, normal_probs_history, epoch, model, args) if saved_memory_normal: del train_queue, valid_queue torch.cuda.empty_cache() count += 1 new_batch_size = args.batch_size + args.batch_increase * count logging.info("new_batch_size = {}".format(new_batch_size)) train_queue = DenseDataLoader(train_dataset, batch_size=new_batch_size, shuffle=True, num_workers=args.batch_size // 2) valid_queue = DenseDataLoader(test_dataset, batch_size=new_batch_size, shuffle=False, num_workers=args.batch_size // 2) # post validation if args.post_val: post_valid_overall_acc, post_valid_class_acc, valid_losses = infer( valid_queue, model, criterion, valid_losses) logging.info('post_valid_overall_acc %f', post_valid_overall_acc) writer.add_scalar('stats/train_acc', train_acc, epoch) writer.add_scalar('stats/valid_overall_acc', valid_overall_acc, epoch) writer.add_scalar('stats/valid_class_acc', valid_class_acc, epoch) utils.save(model, os.path.join(args.save, 'weights.pt')) scheduler.step() logging.info("#" * 30 + " Done " + "#" * 30) logging.info('genotype = %s', model.get_genotype())
if __name__ == "__main__": args = default_argument_parser().parse_args() np.random.seed(args.seed) num_classes = 2 transforms = [] if args.max_points > 0: transforms.append(T.FixedPoints(args.max_points)) if args.augment: transforms.append(T.RandomRotate((-180, 180), axis=2)) # Rotate around z axis transforms.append(T.RandomFlip(0)) # Flp about x axis transforms.append(T.RandomFlip(1)) # Flip about y axis transforms.append(T.RandomTranslate(0.0001)) # Random jitter if args.norm: transforms.append(T.NormalizeScale()) transform = T.Compose(transforms=transforms) if transforms else None train_dataset = EventDataset(args.dataset, "trainval", include_proton=True, task="separation", cleanliness=args.clean, pre_transform=None, transform=transform, balanced_classes=True, fraction=0.001) test_dataset = EventDataset(args.dataset, "test", include_proton=True, task="separation", cleanliness=args.clean,