def minkunet(net_id, pretrained=True, **kwargs):
url_base = 'https://hanlab.mit.edu/files/SPVNAS/minkunet/'
net_config = json.load(open(
download_url(url_base + net_id + '/net.config', model_dir='.torch/minkunet/%s/' % net_id)
))
model = MinkUNet(
num_classes=net_config['num_classes'],
cr=net_config['cr']
).to('cuda:%d'%dist.local_rank() if torch.cuda.is_available() else 'cpu')
if pretrained:
init = torch.load(
download_url(url_base + net_id + '/init', model_dir='.torch/minkunet/%s/' % net_id),
map_location='cuda:%d'%dist.local_rank() if torch.cuda.is_available() else 'cpu'
)['model']
model.load_state_dict(init)
return model
def spvnas_supernet(net_id, pretrained=True, **kwargs):
url_base = 'https://hanlab.mit.edu/files/SPVNAS/spvnas_supernet/'
net_config = json.load(open(
download_url(url_base + net_id + '/net.config', model_dir='.torch/spvnas_supernet/%s/' % net_id)
))
model = SPVNAS(
net_config['num_classes'],
macro_depth_constraint=net_config['macro_depth_constraint'],
pres=net_config['pres'],
vres=net_config['vres']
).to('cuda:%d'%dist.local_rank() if torch.cuda.is_available() else 'cpu')
if pretrained:
init = torch.load(
download_url(url_base + net_id + '/init', model_dir='.torch/spvnas_supernet/%s/' % net_id),
map_location='cuda:%d'%dist.local_rank() if torch.cuda.is_available() else 'cpu'
)['model']
model.load_state_dict(init)
return model
def main():
warnings.filterwarnings("ignore")
# parse args
args, opt = parser.parse_known_args()
opt = parse_unknown_args(opt)
# setup gpu and distributed training
if args.gpu is not None:
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
if not torch.distributed.is_initialized():
dist.init()
torch.backends.cudnn.benchmark = True
torch.cuda.set_device(dist.local_rank())
# setup path
os.makedirs(args.path, exist_ok=True)
# setup random seed
if args.resume:
args.manual_seed = int(time.time())
torch.manual_seed(args.manual_seed)
torch.cuda.manual_seed_all(args.manual_seed)
# load config
exp_config = yaml.safe_load(open(args.config, "r"))
partial_update_config(exp_config, opt)
# save config to run directory
yaml.dump(exp_config,
open(os.path.join(args.path, "config.yaml"), "w"),
sort_keys=False)
# build data_loader
image_size = exp_config["data_provider"]["image_size"]
data_provider, n_classes = build_data_loader(
exp_config["data_provider"]["dataset"],
image_size,
exp_config["data_provider"]["base_batch_size"],
exp_config["data_provider"]["n_worker"],
exp_config["data_provider"]["data_path"],
dist.size(),
dist.rank(),
)
# build model
model = build_model(
exp_config["model"]["name"],
n_classes,
exp_config["model"]["dropout_rate"],
)
print(model)
# netaug
if exp_config.get("netaug", None) is not None:
use_netaug = True
model = augemnt_model(model, exp_config["netaug"], n_classes,
exp_config["model"]["dropout_rate"])
model.set_active(mode="min")
else:
use_netaug = False
# load init
if args.init_from is not None:
init = load_state_dict_from_file(args.init_from)
load_state_dict(model, init, strict=False)
print("Loaded init from %s" % args.init_from)
else:
init_modules(model, init_type=exp_config["run_config"]["init_type"])
print("Random Init")
# profile
profile_model = copy.deepcopy(model)
# during inference, bn will be fused into conv
remove_bn(profile_model)
print(f"Params: {trainable_param_num(profile_model)}M")
print(
f"MACs: {inference_macs(profile_model, data_shape=(1, 3, image_size, image_size))}M"
)
# train
exp_config["generator"] = torch.Generator()
exp_config["generator"].manual_seed(args.manual_seed)
model = nn.parallel.DistributedDataParallel(model.cuda(),
device_ids=[dist.local_rank()])
train(model, data_provider, exp_config, args.path, args.resume, use_netaug)
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')
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(),
])
"tinymbv2",
],
)
parser.add_argument("--init_from", type=str)
parser.add_argument("--reset_bn", action="store_true")
parser.add_argument("--save_path", type=str, default=None)
if __name__ == "__main__":
args = parser.parse_args()
# setup gpu and distributed training
if args.gpu is not None:
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
dist.init()
torch.backends.cudnn.benchmark = True
torch.cuda.set_device(dist.local_rank())
# build data loader
data_loader_dict, n_classes = build_data_loader(
args.dataset,
args.image_size,
args.batch_size,
args.n_worker,
args.data_path,
dist.size(),
dist.rank(),
)
# build model
model = build_model(args.model, n_classes, 0).cuda()