def main() -> None: dist.init() torch.backends.cudnn.benchmark = True torch.cuda.set_device(dist.local_rank()) parser = argparse.ArgumentParser() parser.add_argument('config', metavar='FILE', help='config file') parser.add_argument('--run-dir', metavar='DIR', help='run directory') args, opts = parser.parse_known_args() configs.load(args.config, recursive=True) configs.update(opts) if args.run_dir is None: args.run_dir = auto_set_run_dir() else: set_run_dir(args.run_dir) logger.info(' '.join([sys.executable] + sys.argv)) logger.info(f'Experiment started: "{args.run_dir}".' + '\n' + f'{configs}') dataset = builder.make_dataset() dataflow = {} for split in dataset: sampler = torch.utils.data.DistributedSampler( dataset[split], num_replicas=dist.size(), rank=dist.rank(), shuffle=(split == 'train'), ) dataflow[split] = torch.utils.data.DataLoader( dataset[split], batch_size=configs.batch_size // dist.size(), sampler=sampler, num_workers=configs.workers_per_gpu, pin_memory=True, ) model = builder.make_model() model = torch.nn.parallel.DistributedDataParallel( model.cuda(), device_ids=[dist.local_rank()], ) criterion = builder.make_criterion() optimizer = builder.make_optimizer(model) scheduler = builder.make_scheduler(optimizer) trainer = ClassificationTrainer( model=model, criterion=criterion, optimizer=optimizer, scheduler=scheduler, amp_enabled=configs.amp.enabled, ) trainer.train_with_defaults( dataflow['train'], num_epochs=configs.num_epochs, callbacks=[ SaverRestore(), InferenceRunner( dataflow['test'], callbacks=[ TopKCategoricalAccuracy(k=1, name='acc/top1'), TopKCategoricalAccuracy(k=5, name='acc/top5'), ], ), MaxSaver('acc/top1'), Saver(), ], )
def main() -> None: dist.init() torch.backends.cudnn.benchmark = True torch.cuda.set_device(dist.local_rank()) parser = argparse.ArgumentParser() parser.add_argument('config', metavar='FILE', help='config file') parser.add_argument('--run-dir', metavar='DIR', help='run directory') parser.add_argument('--name', type=str, help='model name') args, opts = parser.parse_known_args() configs.load(args.config, recursive=True) configs.update(opts) if args.run_dir is None: args.run_dir = auto_set_run_dir() else: set_run_dir(args.run_dir) logger.info(' '.join([sys.executable] + sys.argv)) logger.info(f'Experiment started: "{args.run_dir}".' + '\n' + f'{configs}') dataset = builder.make_dataset() dataflow = dict() for split in dataset: sampler = torch.utils.data.distributed.DistributedSampler( dataset[split], num_replicas=dist.size(), rank=dist.rank(), shuffle=(split == 'train')) dataflow[split] = torch.utils.data.DataLoader( dataset[split], batch_size=configs.batch_size if split == 'train' else 1, sampler=sampler, num_workers=configs.workers_per_gpu, pin_memory=True, collate_fn=dataset[split].collate_fn) if 'spvnas' in args.name.lower(): model = spvnas_specialized(args.name) elif 'spvcnn' in args.name.lower(): model = spvcnn(args.name) elif 'mink' in args.name.lower(): model = minkunet(args.name) else: raise NotImplementedError #model = builder.make_model() model = torch.nn.parallel.DistributedDataParallel( model.cuda(), device_ids=[dist.local_rank()], find_unused_parameters=True) model.eval() criterion = builder.make_criterion() optimizer = builder.make_optimizer(model) scheduler = builder.make_scheduler(optimizer) meter = MeanIoU(configs.data.num_classes, configs.data.ignore_label) trainer = SemanticKITTITrainer(model=model, criterion=criterion, optimizer=optimizer, scheduler=scheduler, num_workers=configs.workers_per_gpu, seed=configs.train.seed ) callbacks=Callbacks([ SaverRestore(), MeanIoU( configs.data.num_classes, configs.data.ignore_label ) ]) callbacks._set_trainer(trainer) trainer.callbacks = callbacks trainer.dataflow = dataflow['test'] trainer.before_train() trainer.before_epoch() # important model.eval() for feed_dict in tqdm(dataflow['test'], desc='eval'): _inputs = dict() for key, value in feed_dict.items(): if not 'name' in key: _inputs[key] = value.cuda() inputs = _inputs['lidar'] targets = feed_dict['targets'].F.long().cuda(non_blocking=True) outputs = model(inputs) invs = feed_dict['inverse_map'] all_labels = feed_dict['targets_mapped'] _outputs = [] _targets = [] for idx in range(invs.C[:, -1].max()+1): cur_scene_pts = (inputs.C[:, -1] == idx).cpu().numpy() cur_inv = invs.F[invs.C[:, -1] == idx].cpu().numpy() cur_label = (all_labels.C[:, -1] == idx).cpu().numpy() outputs_mapped = outputs[cur_scene_pts][ cur_inv].argmax(1) targets_mapped = all_labels.F[cur_label] _outputs.append(outputs_mapped) _targets.append(targets_mapped) outputs = torch.cat(_outputs, 0) targets = torch.cat(_targets, 0) output_dict = { 'outputs': outputs, 'targets': targets } trainer.after_step(output_dict) trainer.after_epoch()
def main() -> None: # dist.init() torch.backends.cudnn.benchmark = True # torch.cuda.set_device(dist.local_rank()) parser = argparse.ArgumentParser() parser.add_argument('config', metavar='FILE', help='config file') parser.add_argument('--run-dir', metavar='DIR', help='run directory') parser.add_argument('--pdb', action='store_true', help='pdb') parser.add_argument('--gpu', type=str, help='gpu ids', default=None) args, opts = parser.parse_known_args() configs.load(args.config, recursive=True) configs.update(opts) if configs.debug.pdb or args.pdb: pdb.set_trace() if args.gpu is not None: os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu if configs.debug.set_seed: torch.manual_seed(configs.debug.seed) np.random.seed(configs.debug.seed) if configs.run.device == 'gpu': device = torch.device('cuda') elif configs.run.device == 'cpu': device = torch.device('cpu') else: raise ValueError(configs.run.device) if isinstance(configs.optimizer.lr, str): configs.optimizer.lr = eval(configs.optimizer.lr) if args.run_dir is None: args.run_dir = auto_set_run_dir() else: set_run_dir(args.run_dir) logger.info(' '.join([sys.executable] + sys.argv)) logger.info(f'Experiment started: "{args.run_dir}".' + '\n' + f'{configs}') dataset = builder.make_dataset() dataflow = dict() for split in dataset: sampler = torch.utils.data.RandomSampler(dataset[split]) dataflow[split] = torch.utils.data.DataLoader( dataset[split], batch_size=configs.run.bsz, sampler=sampler, num_workers=configs.run.workers_per_gpu, pin_memory=True) model = builder.make_model() model.to(device) # model = torch.nn.parallel.DistributedDataParallel( # model.cuda(), # device_ids=[dist.local_rank()], # find_unused_parameters=True) total_params = sum(p.numel() for p in model.parameters()) logger.info(f'Model Size: {total_params}') # logger.info(f'Model MACs: {profile_macs(model, inputs)}') criterion = builder.make_criterion() optimizer = builder.make_optimizer(model) scheduler = builder.make_scheduler(optimizer) trainer = LayerRegressionTrainer(model=model, criterion=criterion, optimizer=optimizer, scheduler=scheduler) trainer.train_with_defaults( dataflow['train'], num_epochs=configs.run.n_epochs, callbacks=[ # SaverRestore(), InferenceRunner(dataflow=dataflow['valid'], callbacks=[]), MaxSaver('loss/valid'), # Saver(), ])
def main() -> None: dist.init() torch.backends.cudnn.benchmark = True torch.cuda.set_device(dist.local_rank()) parser = argparse.ArgumentParser() parser.add_argument('config', metavar='FILE', help='config file') parser.add_argument('--run-dir', metavar='DIR', help='run directory') args, opts = parser.parse_known_args() configs.load(args.config, recursive=True) configs.update(opts) if args.run_dir is None: args.run_dir = auto_set_run_dir() else: set_run_dir(args.run_dir) logger.info(' '.join([sys.executable] + sys.argv)) logger.info(f'Experiment started: "{args.run_dir}".' + '\n' + f'{configs}') # seed if ('seed' not in configs.train) or (configs.train.seed is None): configs.train.seed = torch.initial_seed() % (2**32 - 1) seed = configs.train.seed + dist.rank( ) * configs.workers_per_gpu * configs.num_epochs random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed(seed) dataset = builder.make_dataset() dataflow = dict() for split in dataset: sampler = torch.utils.data.distributed.DistributedSampler( dataset[split], num_replicas=dist.size(), rank=dist.rank(), shuffle=(split == 'train')) dataflow[split] = torch.utils.data.DataLoader( dataset[split], batch_size=configs.batch_size, sampler=sampler, num_workers=configs.workers_per_gpu, pin_memory=True, collate_fn=dataset[split].collate_fn) model = builder.make_model() model = torch.nn.parallel.DistributedDataParallel( model.cuda(), device_ids=[dist.local_rank()], find_unused_parameters=True) criterion = builder.make_criterion() optimizer = builder.make_optimizer(model) scheduler = builder.make_scheduler(optimizer) trainer = SemanticKITTITrainer(model=model, criterion=criterion, optimizer=optimizer, scheduler=scheduler, num_workers=configs.workers_per_gpu, seed=seed) trainer.train_with_defaults( dataflow['train'], num_epochs=configs.num_epochs, callbacks=[ InferenceRunner(dataflow[split], callbacks=[ MeanIoU(name=f'iou/{split}', num_classes=configs.data.num_classes, ignore_label=configs.data.ignore_label) ]) for split in ['test'] ] + [ MaxSaver('iou/test'), Saver(), ])