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 get_components(cfg):
labeled_loader, unlabeled_loader, val_loader, test_loader = get_dataloader(cfg.data.name,
cfg.data.labeled_size,
cfg.data.unlabeled_size,
cfg.data.val_size,
cfg.data.batch_size,
cfg.data.random_state,
download=cfg.data.download,
pilaugment=cfg.data.get('pilaugment',
False)
)
model = wrn28_2(num_classes=6 if cfg.data.name == "animal" else 10)
optimizer = {'adam': optim.Adam(lr=cfg.optim.lr),
'sgd': optim.SGD(lr=cfg.optim.lr, momentum=0.9)}[cfg.optim.name]
scheduler = {'adam': None,
'sgd': lr_scheduler.CosineAnnealingWithWarmup(cfg.optim.epochs,
4, cfg.optim.epochs // 100)}[cfg.optim.name]
ema_model = partial(EMAModel, ema_rate=cfg.model.ema_rate, weight_decay=cfg.optim.wd * cfg.optim.lr)
num_classes = {"animal": 6, "cifar100": 100, "tinyimagenet": 200}.get(cfg.data.name, 10)
tq = reporters.TQDMReporter(range(cfg.optim.epochs))
_callbacks = [callbacks.AccuracyCallback(),
callbacks.LossCallback(),
reporters.IOReporter("."),
reporters.TensorboardReporter("."), tq]
return PackedLoader(labeled_loader, unlabeled_loader), val_loader, test_loader, model, optimizer, \
scheduler, ema_model, num_classes, tq, _callbacks
def train_and_eval(cfg):
train_loader, val_loader, test_loader, num_classes = get_dataloader(
cfg.data.name, cfg.data.val_size, cfg.data.batch_size,
cfg.data.download, cfg.augment, False)
model = get_model(cfg.model.name, num_classes)
optimizer = optim.SGD(cfg.optim.model.lr,
momentum=0.9,
weight_decay=cfg.optim.model.weight_decay)
scheduler = lr_scheduler.MultiStepLR(cfg.optim.model.steps)
tq = reporters.TQDMReporter(range(cfg.optim.epochs), verb=cfg.verb)
callback = [
callbacks.AccuracyCallback(),
callbacks.LossCallback(),
reporters.TensorboardReporter("."),
reporters.IOReporter("."), tq
]
with trainers.SupervisedTrainer(model,
optimizer,
F.cross_entropy,
callbacks=callback,
scheduler=scheduler) as trainer:
for ep in tq:
trainer.train(train_loader)
trainer.test(val_loader, 'val')
trainer.test(test_loader)
def main(cfg):
model = wrn28_2(num_classes=10)
train_loader, test_loader = get_dataloaders(cfg.data.name,
cfg.data.batch_size,
cfg.data.train_size,
cfg.data.random_state)
optimizer = optim.Adam(lr=cfg.optim.lr)
tq = reporters.TQDMReporter(range(cfg.optim.epochs))
c = [
callbacks.AccuracyCallback(),
callbacks.LossCallback(),
reporters.IOReporter("."), tq
]
with SupervisedTrainer(
model,
optimizer,
F.cross_entropy,
callbacks=c,
ema_model=partial(EMAModel,
ema_rate=cfg.model.ema_rate,
weight_decay=cfg.optim.wd * cfg.optim.lr),
) as trainer:
for _ in tq:
trainer.train(train_loader)
trainer.test(test_loader)
trainer.logger.info(
f"test accuracy: {median(c[0].history['test'][-20:])}")
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 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)
