def main():
if not torch.cuda.is_available():
logging.info('no gpu device available')
sys.exit(1)
if args.auxiliary and args.net_type == 'macro':
logging.info('auxiliary head classifier not supported for macro search space models')
sys.exit(1)
logging.info("args = %s", args)
cudnn.enabled = True
cudnn.benchmark = True
np.random.seed(args.seed)
torch.manual_seed(args.seed)
torch.cuda.manual_seed(args.seed)
# Data
_, valid_transform = utils._data_transforms_cifar10(args)
valid_data = torchvision.datasets.CIFAR10(root=args.data, train=False, download=True, transform=valid_transform)
valid_queue = torch.utils.data.DataLoader(
valid_data, batch_size=args.batch_size, shuffle=False, pin_memory=True, num_workers=1)
# classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')
# Model
if args.net_type == 'micro':
logging.info("==> Building micro search space encoded architectures")
genotype = eval("genotypes.%s" % args.arch)
net = PyrmNASNet(args.init_channels, num_classes=10, layers=args.layers,
auxiliary=args.auxiliary, genotype=genotype,
increment=args.filter_increment, SE=args.SE)
elif args.net_type == 'macro':
genome = eval("macro_genotypes.%s" % args.arch)
channels = [(3, 128), (128, 128), (128, 128)]
net = EvoNetwork(genome, channels, 10, (32, 32), decoder='dense')
else:
raise NameError('Unknown network type, please only use supported network type')
# logging.info("{}".format(net))
logging.info("param size = %fMB", utils.count_parameters_in_MB(net))
net = net.to(device)
# no drop path during inference
net.droprate = 0.0
utils.load(net, args.model_path)
criterion = nn.CrossEntropyLoss()
criterion.to(device)
# inference on original CIFAR-10 test images
infer(valid_queue, net, criterion)
def main():
if not torch.cuda.is_available():
logging.info('no gpu device available')
sys.exit(1)
if args.auxiliary and args.net_type == 'macro':
logging.info(
'auxiliary head classifier not supported for macro search space models'
)
sys.exit(1)
logging.info("args = %s", args)
cudnn.enabled = True
cudnn.benchmark = True
np.random.seed(args.seed)
torch.manual_seed(args.seed)
torch.cuda.manual_seed(args.seed)
best_acc = 0 # initiate a artificial best accuracy so far
# Data
train_transform, valid_transform = utils._data_transforms_cifar10(args)
train_data = torchvision.datasets.CIFAR10(root=args.data,
train=True,
download=True,
transform=train_transform)
valid_data = torchvision.datasets.CIFAR10(root=args.data,
train=False,
download=True,
transform=valid_transform)
train_queue = torch.utils.data.DataLoader(train_data,
batch_size=args.batch_size,
shuffle=True,
pin_memory=True,
num_workers=2)
valid_queue = torch.utils.data.DataLoader(valid_data,
batch_size=128,
shuffle=False,
pin_memory=True,
num_workers=2)
# classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')
# Model
if args.net_type == 'micro':
logging.info("==> Building micro search space encoded architectures")
genotype = eval("genotypes.%s" % args.arch)
net = PyrmNASNet(args.init_channels,
num_classes=10,
layers=args.layers,
auxiliary=args.auxiliary,
genotype=genotype,
increment=args.filter_increment,
SE=args.SE)
elif args.net_type == 'macro':
genome = eval("macro_genotypes.%s" % args.arch)
channels = [(3, 128), (128, 128), (128, 128)]
net = EvoNetwork(genome, channels, 10, (32, 32), decoder='dense')
else:
raise NameError(
'Unknown network type, please only use supported network type')
# logging.info("{}".format(net))
logging.info("param size = %fMB", utils.count_parameters_in_MB(net))
net = net.to(device)
n_epochs = args.epochs
parameters = filter(lambda p: p.requires_grad, net.parameters())
criterion = nn.CrossEntropyLoss()
criterion.to(device)
optimizer = optim.SGD(parameters,
lr=args.learning_rate,
momentum=args.momentum,
weight_decay=args.weight_decay)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
optimizer, n_epochs, eta_min=args.min_learning_rate)
for epoch in range(n_epochs):
scheduler.step()
logging.info('epoch %d lr %e', epoch, scheduler.get_lr()[0])
net.droprate = args.droprate * epoch / args.epochs
train(train_queue, net, criterion, optimizer)
_, valid_acc = infer(valid_queue, net, criterion)
if valid_acc > best_acc:
utils.save(net, os.path.join(args.save, 'weights.pt'))
best_acc = valid_acc