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 main():
train_loader, test_loader = imagenet_loaders(
args.root,
args.batch_size,
num_train_samples=args.batch_size * args.train_max_iter,
num_test_samples=args.batch_size * args.test_max_iter)
pretrained_model = resnet50(pretrained=True)
for p in pretrained_model.parameters():
p.requires_grad = False
pretrained_model.eval()
generator = ResNetGenerator(3, 3, args.num_filters)
generator.cuda()
optimizer = optim.Adam(lr=args.lr, betas=(args.beta1, 0.999))
trainer = Trainer({
"generator": generator,
"classifier": pretrained_model
},
optimizer,
reporter.TensorboardReporter([
adv_accuracy, fooling_rate,
callbacks.AccuracyCallback(),
callbacks.LossCallback()
],
save_dir="results"),
noise=torch.randn(3, 224,
224).expand(args.batch_size, -1, -1,
-1))
for ep in range(args.epochs):
trainer.train(train_loader)
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 search(cfg: BaseConfig):
train_loader, _, num_classes = DATASET_REGISTRY(cfg.data.name)(
batch_size=cfg.data.batch_size,
train_size=cfg.data.train_size,
drop_last=True,
download=cfg.data.download,
return_num_classes=True,
num_workers=4)
model = {
'main':
Discriminator(MODEL_REGISTRY('wrn40_2')(num_classes)),
'policy':
Policy.faster_auto_augment_policy(cfg.model.num_sub_policies,
cfg.model.temperature,
cfg.model.operation_count,
cfg.model.num_chunks)
}
optimizer = {
'main': optim.Adam(lr=cfg.optim.main_lr, betas=(0, 0.999)),
'policy': optim.Adam(lr=cfg.optim.policy_lr, betas=(0, 0.999))
}
tqdm = callbacks.TQDMReporter(range(cfg.optim.epochs))
c = [
callbacks.LossCallback(), # classification loss
callbacks.metric_callback_by_name('d_loss'), # discriminator loss
callbacks.metric_callback_by_name('a_loss'), # augmentation loss
tqdm
]
with AdvTrainer(model,
optimizer,
F.cross_entropy,
callbacks=c,
cfg=cfg.model,
use_cuda_nonblocking=True) as trainer:
for _ in tqdm:
trainer.train(train_loader)
trainer.save(
pathlib.Path(hydra.utils.get_original_cwd()) / 'policy_weights' /
cfg.data.name)
if __name__ == '__main__':
import miniargs
from torch.nn import functional as F
p = miniargs.ArgumentParser()
p.add_int("--batch_size", default=128)
p.add_int("--epochs", default=300)
p.add_str("--optimizer", choices=["sgd", "adam"])
p.add_float("--lr", default=1e-2)
p.add_multi_str("--group", default=["conv1", "layer1", "layer2", "layer3"])
p.add_int("--step", default=50)
p.add_int("--num_convs", default=3)
p.add_int("--num_fcs", default=3)
args = p.parse()
optimizer = {"adam": optim.Adam(lr=3e-4, weight_decay=1e-4),
"sgd": optim.SGD(lr=args.lr, momentum=0.9, weight_decay=1e-4)}[args.optimizer]
train_loader, test_loader = cifar10_loaders(args.batch_size)
resnet = module_converter(resnet56(num_classes=10), keys=["conv1", "bn1", "relu", "layer1", "layer2", "layer3"])
aux = nn.ModuleDict(OrderedDict({k: v for k, v in {
# 32x32
"conv1": generate_aux(32, 16, 10, args.num_convs, args.num_fcs),
# 32x32
"layer1": generate_aux(32, 16, 10, args.num_convs, args.num_fcs),
# 16x16
"layer2": generate_aux(16, 32, 10, args.num_convs, args.num_fcs),
# 8x8
"layer3": generate_aux(8, 64, 10, args.num_convs, args.num_fcs),
}.items() if k in args.group}))
model = NaiveGreedyModule(resnet, aux=aux,