def main():
if not torch.cuda.is_available():
logging.info('no gpu device available')
sys.exit(1)
np.random.seed(args.seed)
torch.cuda.set_device(args.gpu)
cudnn.benchmark = True
torch.manual_seed(args.seed)
cudnn.enabled = True
torch.cuda.manual_seed(args.seed)
logging.info('gpu device = %d' % args.gpu)
logging.info("args = %s", args)
genotype = eval("genotypes.%s" % args.arch)
model = Network(args.init_channels, CIFAR_CLASSES, args.layers,
args.auxiliary, genotype)
model = model.cuda()
utils.load(model, args.model_path)
logging.info("param size = %fMB", utils.count_parameters_in_MB(model))
criterion = nn.CrossEntropyLoss()
criterion = criterion.cuda()
_, test_transform = utils._data_transforms_cifar10(args)
test_data = dset.CIFAR10(root=args.data,
train=False,
download=True,
transform=test_transform)
test_queue = torch.utils.data.DataLoader(test_data,
batch_size=args.batch_size,
shuffle=False,
pin_memory=True,
num_workers=2)
model.drop_path_prob = args.drop_path_prob
test_acc, test_obj = infer(test_queue, model, criterion)
logging.info('test_acc %f', test_acc)
def load(self, epoch=None):
if epoch is not None:
model_obj_path = os.path.join(self.args.save, 'one_shot_model_{}.obj'.format(epoch))
if os.path.exists(model_obj_path):
utils.load(self.model, model_obj_path)
else:
model_pt_path = os.path.join(self.args.save, 'one_shot_model_{}.pt'.format(epoch))
utils.load(self.model, model_pt_path)
else:
utils.load(self.model, os.path.join(self.args.save, 'weights.obj'))
def load(self, epoch=None, ft=False):
if ft:
model_pt_path = os.path.join('/home/ubuntu/nasbench-1shot1_crwhite/experiments/', 'ft_model_{}.pt'.format(epoch))
utils.load(self.model, model_pt_path)
print('loaded model', model_pt_path)
else:
if epoch is not None:
model_obj_path = os.path.join(self.args.save, 'one_shot_model_{}.obj'.format(epoch))
if os.path.exists(model_obj_path):
utils.load(self.model, model_obj_path)
else:
model_pt_path = os.path.join(self.args.save, 'one_shot_model_{}.pt'.format(epoch))
utils.load(self.model, model_pt_path)
print('loaded model', model_pt_path)
else:
utils.load(self.model, os.path.join(self.args.save, 'weights.obj'))