def main(): if args.distributed: init_distributed() model = se_resnet50(num_classes=1000) optimizer = optim.SGD(lr=0.6 / 1024 * args.batch_size, momentum=0.9, weight_decay=1e-4) scheduler = lr_scheduler.MultiStepLR([50, 70]) train_loader, test_loader = imagenet_loaders(args.root, args.batch_size, distributed=args.distributed, num_train_samples=args.batch_size * 10 if args.debug else None, num_test_samples=args.batch_size * 10 if args.debug else None) c = [callbacks.AccuracyCallback(), callbacks.AccuracyCallback(k=5), callbacks.LossCallback(), callbacks.WeightSave('.'), reporters.TensorboardReporter('.'), reporters.TQDMReporter(range(args.epochs))] with SupervisedTrainer(model, optimizer, F.cross_entropy, callbacks=c, scheduler=scheduler, ) as trainer: for _ in c[-1]: trainer.train(train_loader) trainer.test(test_loader)
def main(): model = { "resnet20": resnet20, "wrn28_10": wrn28_10 }[args.model](num_classes=10) weight_decay = {"resnet20": 1e-4, "wrn28_10": 5e-4}[args.model] lr_decay = {"resnet20": 0.1, "wrn28_10": 0.2}[args.model] train_loader, test_loader = cifar10_loaders(args.batch_size) optimizer = optim.SGD(lr=1e-1, momentum=0.9, weight_decay=weight_decay) scheduler = lr_scheduler.MultiStepLR([100, 150], gamma=lr_decay) tq = reporters.TQDMReporter(range(args.epochs), verb=True) c = [ callbacks.AccuracyCallback(), callbacks.LossCallback(), reporters.IOReporter("."), reporters.TensorboardReporter("."), callbacks.WeightSave("."), tq ] with trainers.SupervisedTrainer(model, optimizer, F.cross_entropy, callbacks=c, scheduler=scheduler) as trainer: for _ in tq: trainer.train(train_loader) trainer.test(test_loader)
def main(): model = se_resnet50(num_classes=1000) optimizer = optim.SGD(lr=0.6 / 1024 * args.batch_size, momentum=0.9, weight_decay=1e-4) scheduler = lr_scheduler.MultiStepLR([50, 70]) c = [callbacks.AccuracyCallback(), callbacks.LossCallback()] r = reporters.TQDMReporter(range(args.epochs), callbacks=c) tb = reporters.TensorboardReporter(c) rep = callbacks.CallbackList(r, tb, callbacks.WeightSave("checkpoints")) if args.distributed: # DistributedSupervisedTrainer sets up torch.distributed if args.local_rank == 0: print("\nuse DistributedDataParallel") trainer = DistributedSupervisedTrainer(model, optimizer, F.cross_entropy, callbacks=rep, scheduler=scheduler, init_method=args.init_method, backend=args.backend) else: multi_gpus = torch.cuda.device_count() > 1 if multi_gpus: print("\nuse DataParallel") trainer = SupervisedTrainer(model, optimizer, F.cross_entropy, callbacks=rep, scheduler=scheduler, data_parallel=multi_gpus) # if distributed, need to setup loaders after DistributedSupervisedTrainer train_loader, test_loader = imagenet_loaders(args.root, args.batch_size, distributed=args.distributed, num_train_samples=args.batch_size * 10 if args.debug else None, num_test_samples=args.batch_size * 10 if args.debug else None) for _ in r: trainer.train(train_loader) trainer.test(test_loader)
def main(): if is_distributed(): init_distributed() model = se_resnet50(num_classes=1000) optimizer = optim.SGD(lr=0.6 / 1024 * args.batch_size, momentum=0.9, weight_decay=1e-4) scheduler = lr_scheduler.MultiStepLR([50, 70]) train_loader, test_loader = DATASET_REGISTRY("imagenet")(args.batch_size) c = [ callbacks.AccuracyCallback(), callbacks.AccuracyCallback(k=5), callbacks.LossCallback(), callbacks.WeightSave("."), reporters.TensorboardReporter("."), reporters.TQDMReporter(range(args.epochs)), ] with SupervisedTrainer( model, optimizer, F.cross_entropy, callbacks=c, scheduler=scheduler, ) as trainer: for _ in c[-1]: trainer.train(train_loader) trainer.test(test_loader)
def main(): if args.distributed: init_distributed() if args.enable_accimage: enable_accimage() model = resnet50() optimizer = optim.SGD(lr=1e-1 * args.batch_size * get_num_nodes() / 256, momentum=0.9, weight_decay=1e-4) scheduler = lr_scheduler.MultiStepLR([30, 60, 80]) c = [callbacks.AccuracyCallback(), callbacks.LossCallback()] r = reporters.TQDMReporter(range(args.epochs), callbacks=c) tb = reporters.TensorboardReporter(c) rep = callbacks.CallbackList(r, tb, callbacks.WeightSave("checkpoints")) _train_loader, _test_loader = imagenet_loaders( args.root, args.batch_size, distributed=args.distributed, num_train_samples=args.batch_size * 10 if args.debug else None, num_test_samples=args.batch_size * 10 if args.debug else None) if args.distributed: # DistributedSupervisedTrainer sets up torch.distributed if args.local_rank == 0: print("\nuse DistributedDataParallel\n") trainer = DistributedSupervisedTrainer(model, optimizer, F.cross_entropy, callbacks=rep, scheduler=scheduler, init_method=args.init_method, backend=args.backend, enable_amp=args.enable_amp) else: use_multi_gpus = torch.cuda.device_count() > 1 if use_multi_gpus: print("\nuse DataParallel\n") trainer = SupervisedTrainer(model, optimizer, F.cross_entropy, callbacks=rep, data_parallel=use_multi_gpus) for epoch in r: if args.use_prefetcher: train_loader = prefetcher.DataPrefetcher(_train_loader) test_loader = prefetcher.DataPrefetcher(_test_loader) else: train_loader, test_loader = _train_loader, _test_loader # following apex's training scheme trainer.train(train_loader) trainer.test(test_loader) rep.close()
def main(cfg): model = { "resnet20": resnet20, "wrn28_10": wrn28_10 }[cfg.model](num_classes=10) weight_decay = {"resnet20": 1e-4, "wrn28_10": 5e-4}[cfg.model] lr_decay = {"resnet20": 0.1, "wrn28_10": 0.2}[cfg.model] train_loader, test_loader = vision_loaders("cifar10", cfg.batch_size) optimizer = None if cfg.bn_no_wd else optim.SGD( lr=1e-1, momentum=0.9, weight_decay=weight_decay) scheduler = lr_scheduler.MultiStepLR([100, 150], gamma=lr_decay) tq = reporters.TQDMReporter(range(cfg.epochs), verb=True) c = [ callbacks.AccuracyCallback(), callbacks.LossCallback(), reporters.IOReporter("."), reporters.TensorboardReporter("."), callbacks.WeightSave("."), tq ] if cfg.bn_no_wd: def set_optimizer(trainer): bn_params = [] non_bn_parameters = [] for name, p in trainer.model.named_parameters(): if "bn" in name: bn_params.append(p) else: non_bn_parameters.append(p) optim_params = [ { "params": bn_params, "weight_decay": 0 }, { "params": non_bn_parameters, "weight_decay": weight_decay }, ] trainer.optimizer = torch.optim.SGD(optim_params, lr=1e-1, momentum=0.9) trainers.SupervisedTrainer.set_optimizer = set_optimizer with trainers.SupervisedTrainer(model, optimizer, F.cross_entropy, callbacks=c, scheduler=scheduler) as trainer: for _ in tq: trainer.train(train_loader) trainer.test(test_loader)
def test(tmp_path, rep, save_freq): temp_dir = tmp_path / "test" @callbacks.metric_callback_decorator def ca(data): output, target = data["output"], data["data"][1] return { i: v for i, v in enumerate(metrics.classwise_accuracy(output, target)) } model = nn.Linear(10, 10) optimizer = optim.SGD(lr=0.1) c = callbacks.CallbackList( callbacks.AccuracyCallback(), ca, callbacks.WeightSave(save_path=temp_dir, save_freq=save_freq)) epoch = range(1) loader = [(torch.randn(2, 10), torch.zeros(2, dtype=torch.long)) for _ in range(10)] with { "tqdm": lambda: reporters.TQDMReporter(epoch, c, temp_dir), "logger": lambda: reporters.LoggerReporter(c, temp_dir), "tensorboard": lambda: reporters.TensorboardReporter(c, temp_dir) }[rep]() as _rep: tr = trainers.SupervisedTrainer(model, optimizer, F.cross_entropy, callbacks=_rep, verb=False) if rep == "tqdm": epoch = _rep for _ in epoch: tr.train(loader) tr.test(loader) tr.exit() try: # .../test/**/0.pkl save_file = list(Path(temp_dir).glob("*/*.pkl"))[0] except IndexError as e: print(list(Path(temp_dir).glob("*/*"))) raise e tr.resume(save_file) c = callbacks.AccuracyCallback() with { "tqdm": lambda: reporters.TQDMReporter(epoch, c, temp_dir), "logger": lambda: reporters.LoggerReporter(c, temp_dir), "tensorboard": lambda: reporters.TensorboardReporter(c, temp_dir) }[rep]() as _rep: inferencer = Inferencer(model, _rep) inferencer.load(save_file) inferencer.run(loader)
def test(rep): tmpdir = str(gettempdir()) if rep == "tensorboard" and not is_tensorboardX_available: pytest.skip("tensorboardX is not available") @callbacks.metric_callback_decorator def ca(data): output, target = data["output"], data["data"][1] return { i: v for i, v in enumerate(metrics.classwise_accuracy(output, target)) } model = nn.Linear(10, 10) optimizer = optim.SGD(lr=0.1) c = callbacks.CallbackList(callbacks.AccuracyCallback(), ca, callbacks.WeightSave(tmpdir)) epoch = range(1) loader = [(torch.randn(2, 10), torch.zeros(2, dtype=torch.long)) for _ in range(10)] with { "tqdm": lambda: reporters.TQDMReporter(epoch, c, tmpdir), "logger": lambda: reporters.LoggerReporter(c, tmpdir), "tensorboard": lambda: reporters.TensorboardReporter(c, tmpdir) }[rep]() as _rep: tr = trainers.SupervisedTrainer(model, optimizer, F.cross_entropy, callbacks=_rep, verb=False) if rep == "tqdm": epoch = _rep for _ in epoch: tr.train(loader) tr.test(loader) save_file = list(Path(tmpdir).glob("*/*.pkl"))[0] tr.resume(save_file) c = callbacks.AccuracyCallback() with { "tqdm": lambda: reporters.TQDMReporter(epoch, c, tmpdir), "logger": lambda: reporters.LoggerReporter(c, tmpdir), "tensorboard": lambda: reporters.TensorboardReporter(c, tmpdir) }[rep]() as _rep: inferencer = Inferencer(model, _rep) inferencer.load(save_file) inferencer.run(loader)
def main(): train_loader, test_loader = get_dataloader(args.batch_size, args.root) gpus = list(range(torch.cuda.device_count())) se_resnet = nn.DataParallel(se_resnet50(num_classes=1000), device_ids=gpus) optimizer = optim.SGD(lr=0.6 / 1024 * args.batch_size, momentum=0.9, weight_decay=1e-4) scheduler = lr_scheduler.StepLR(30, gamma=0.1) weight_saver = callbacks.WeightSave("checkpoints") tqdm_rep = reporter.TQDMReporter(range(args.epochs), callbacks=[callbacks.AccuracyCallback()]) trainer = Trainer(se_resnet, optimizer, F.cross_entropy, scheduler=scheduler, callbacks=callbacks.CallbackList(weight_saver, tqdm_rep)) for _ in tqdm_rep: trainer.train(train_loader) trainer.test(test_loader)
def main(batch_size): layers = ["layer1.0.conv1", "layer2.0.conv1", "layer3.0.conv1", "fc"] train_loader, test_loader = cifar10_loaders(128) weight_save = callbacks.WeightSave("checkpoints") model = resnet20(num_classes=10) model2 = deepcopy(model) optimizer = torch.optim.SGD(params=model.parameters(), lr=0.1, momentum=0.9, weight_decay=1e-4) scheduler = torch.optim.lr_scheduler.StepLR(optimizer, 50) trainer = trainers.SupervisedTrainer(model, optimizer, F.cross_entropy, scheduler=scheduler, callbacks=weight_save, verb=False) for ep in trange(100, ncols=80): trainer.train(train_loader) hooks1 = [CCAHook(model, name, svd_device=args.device) for name in layers] hooks2 = [CCAHook(model2, name, svd_device=args.device) for name in layers] device = next(model.parameters()).device model2.to(device) input = hooks1[0].data(train_loader.dataset, batch_size=batch_size).to(device) history = [] def distance(): model.eval() model2.eval() with torch.no_grad(): model(input) model2(input) return [h1.distance(h2) for h1, h2 in zip(hooks1, hooks2)] # 0 and 99 history.append(distance()) # 29 and 99, ... for ep in (29, 49, 99): saved = torch.load(weight_save.save_path / f"{ep}.pkl") model2.load_state_dict(saved["model"]) history.append(distance()) plot(history, layers)