def main(cfg: Config):
if cfg.enable_accimage:
enable_accimage()
model = resnet50()
optimizer = optim.SGD(lr=1e-1 * cfg.batch_size * get_num_nodes() / 256, momentum=0.9, weight_decay=1e-4)
scheduler = lr_scheduler.MultiStepLR([30, 60, 80])
train_loader, test_loader = DATASET_REGISTRY("fast_imagenet" if cfg.use_fast_collate else
"imagenet")(cfg.batch_size,
train_size=cfg.batch_size * 50 if cfg.debug else None,
test_size=cfg.batch_size * 50 if cfg.debug else None,
num_workers=cfg.num_workers)
use_multi_gpus = not is_distributed() and torch.cuda.device_count() > 1
with SupervisedTrainer(model,
optimizer,
F.cross_entropy,
reporters=[reporters.TensorboardReporter(".")],
scheduler=scheduler,
data_parallel=use_multi_gpus,
use_amp=cfg.use_amp,
use_cuda_nonblocking=True,
use_sync_bn=cfg.use_sync_bn,
report_accuracy_topk=5) as trainer:
for epoch in trainer.epoch_range(cfg.epochs):
trainer.train(train_loader)
trainer.test(test_loader)
print(f"Max Test Accuracy={max(trainer.reporter.history('accuracy/test')):.3f}")
def main(cfg):
if cfg.use_accimage:
enable_accimage()
model = MODEL_REGISTRY(cfg.model.name)(num_classes=10)
train_loader, test_loader = DATASET_REGISTRY("fast_cifar10" if cfg.use_fast_collate else "cifar10"
)(cfg.data.batch_size, num_workers=4,
use_prefetcher=cfg.use_prefetcher)
optimizer = None if cfg.bn_no_wd else optim.SGD(lr=1e-1, momentum=0.9, weight_decay=cfg.optim.weight_decay)
scheduler = lr_scheduler.MultiStepLR([100, 150], gamma=cfg.optim.lr_decay)
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": cfg.optim.weight_decay},
]
trainer.optimizer = torch.optim.SGD(optim_params, lr=1e-1, momentum=0.9)
trainers.SupervisedTrainer.set_optimizer = set_optimizer
if cfg.use_zerograd_none:
import types
def set_optimizer(trainer):
# see Apex for details
def zero_grad(self):
for group in self.param_groups:
for p in group['params']:
p.grad = None
trainer.optimizer = trainer.optimizer(trainer.model.parameters())
trainer.optimizer.zero_grad = types.MethodType(zero_grad, trainer.optimizer)
trainers.SupervisedTrainer.set_optimizer = set_optimizer
with trainers.SupervisedTrainer(model,
optimizer,
F.cross_entropy,
reporters=[reporters.TensorboardReporter('.')],
scheduler=scheduler,
use_amp=cfg.use_amp,
debug=cfg.debug
) as trainer:
for _ in trainer.epoch_range(cfg.optim.epochs):
trainer.train(train_loader)
trainer.test(test_loader)
print(f"Max Test Accuracy={max(trainer.reporter.history('accuracy/test')):.3f}")
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):
if cfg.use_accimage:
enable_accimage()
model = MODEL_REGISTRY(cfg.name)(num_classes=10)
train_loader, test_loader = DATASET_REGISTRY(
"fast_cifar10" if cfg.use_fast_collate else "cifar10")(
cfg.batch_size, num_workers=4, use_prefetcher=cfg.use_prefetcher)
optimizer = None if cfg.bn_no_wd else optim.SGD(
lr=cfg.lr, momentum=0.9, weight_decay=cfg.weight_decay)
scheduler = lr_scheduler.CosineAnnealingWithWarmup(cfg.epochs, 4, 5)
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": cfg.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,
reporters=[reporters.TensorboardReporter('.')],
scheduler=scheduler,
use_amp=cfg.use_amp,
debug=cfg.debug) as trainer:
for _ in trainer.epoch_range(cfg.epochs):
trainer.train(train_loader)
trainer.test(test_loader)
trainer.scheduler.step()
print(
f"Max Test Accuracy={max(trainer.reporter.history('accuracy/test')):.3f}"
)
def main(cfg):
if cfg.distributed.enable:
init_distributed(use_horovod=cfg.distributed.use_horovod,
backend=cfg.distributed.backend,
init_method=cfg.distributed.init_method)
if cfg.enable_accimage:
enable_accimage()
model = resnet50()
optimizer = optim.SGD(lr=1e-1 * cfg.batch_size * get_num_nodes() / 256,
momentum=0.9,
weight_decay=1e-4)
scheduler = lr_scheduler.MultiStepLR([30, 60, 80])
tq = reporters.TQDMReporter(range(cfg.epochs))
c = [
callbacks.AccuracyCallback(),
callbacks.AccuracyCallback(k=5),
callbacks.LossCallback(), tq,
reporters.TensorboardReporter("."),
reporters.IOReporter(".")
]
_train_loader, _test_loader = imagenet_loaders(
cfg.root,
cfg.batch_size,
distributed=cfg.distributed.enable,
num_train_samples=cfg.batch_size * 10 if cfg.debug else None,
num_test_samples=cfg.batch_size * 10 if cfg.debug else None)
use_multi_gpus = not cfg.distributed.enable and torch.cuda.device_count(
) > 1
with SupervisedTrainer(model,
optimizer,
F.cross_entropy,
callbacks=c,
scheduler=scheduler,
data_parallel=use_multi_gpus,
use_horovod=cfg.distributed.use_horovod) as trainer:
for epoch in tq:
if cfg.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)
def main(cfg):
if cfg.use_accimage:
enable_accimage()
data = DATASET_REGISTRY(cfg.data).setup(
cfg.batch_size,
num_workers=4,
download=cfg.download,
prefetch_factor=cfg.prefetch_factor,
persistent_workers=cfg.persistent_workers)
model = MODEL_REGISTRY(cfg.model)(num_classes=data.num_classes)
optimizer = None if cfg.bn_no_wd else optim.SGD(
lr=cfg.lr,
momentum=0.9,
weight_decay=cfg.weight_decay,
multi_tensor=cfg.use_multi_tensor)
scheduler = lr_scheduler.CosineAnnealingWithWarmup(cfg.epochs, 4, 5)
if cfg.bn_no_wd:
def set_optimizer(trainer):
bn_params = []
non_bn_parameters = []
for name, p in trainer.model.named_parameters():
if "norm" 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": cfg.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,
reporters=[reporters.TensorboardReporter('.')],
scheduler=scheduler,
use_amp=cfg.use_amp,
use_channel_last=cfg.use_channel_last,
debug=cfg.debug) as trainer:
for _ in trainer.epoch_range(cfg.epochs):
trainer.train(data.train_loader)
trainer.test(data.test_loader)
trainer.scheduler.step()
print(
f"Max Test Accuracy={max(trainer.reporter.history('accuracy/test')):.3f}"
)