def setup_model(opa=False, dataset='imagenet', model_size='SMALL'):
seed = 42
restart_from = 50
n_epochs = 100
pretrained = False
n_gpus = 1
use_group_norm = False
shine = False
fpn = False
adjoint_broyden = True
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
try:
torch.multiprocessing.set_start_method('spawn')
except RuntimeError:
pass
args = update_config_w_args(
n_epochs=n_epochs,
pretrained=pretrained,
n_gpus=n_gpus,
dataset=dataset,
model_size=model_size,
use_group_norm=use_group_norm,
)
torch.set_default_tensor_type('torch.cuda.FloatTensor')
logger, final_output_dir, tb_log_dir = create_logger(
config,
args.cfg,
'train',
shine=shine,
fpn=fpn,
seed=seed,
use_group_norm=use_group_norm,
adjoint_broyden=adjoint_broyden,
opa=True,
)
model = eval('models.' + config.MODEL.NAME + '.get_cls_net')(
config,
shine=shine,
fpn=fpn,
gradient_correl=False,
gradient_ratio=False,
adjoint_broyden=adjoint_broyden,
opa=opa,
).cuda()
resume_file = f'checkpoint_{restart_from}.pth.tar'
model_state_file = os.path.join(final_output_dir, resume_file)
checkpoint = torch.load(model_state_file)
model.load_state_dict(checkpoint['state_dict'])
return model
def evaluate_classifier(
n_gpus=1,
dataset='imagenet',
model_size='SMALL',
shine=False,
fpn=False,
refine=False,
n_refine=None,
fallback=False,
n_samples=None,
seed=0,
adjoint_broyden=False,
opa=False,
use_group_norm=False,
check_contract=False,
n_iter=20,
):
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
args = update_config_w_args(
n_gpus=n_gpus,
dataset=dataset,
model_size=model_size,
use_group_norm=use_group_norm,
)
logger, final_output_dir, tb_log_dir = create_logger(
config,
args.cfg,
'valid',
shine=shine,
fpn=fpn,
seed=seed,
use_group_norm=use_group_norm,
adjoint_broyden=adjoint_broyden,
opa=opa,
refine=refine,
n_refine=n_refine,
fallback=fallback,
)
logger.info(pprint.pformat(args))
logger.info(pprint.pformat(config))
# cudnn related setting
cudnn.benchmark = config.CUDNN.BENCHMARK
torch.backends.cudnn.deterministic = config.CUDNN.DETERMINISTIC
torch.backends.cudnn.enabled = config.CUDNN.ENABLED
model = eval('models.' + config.MODEL.NAME + '.get_cls_net')(
config,
shine=shine,
fpn=fpn,
refine=refine,
fallback=fallback,
adjoint_broyden=adjoint_broyden)
dump_input = torch.rand(
(1, 3, config.MODEL.IMAGE_SIZE[1], config.MODEL.IMAGE_SIZE[0]))
logger.info(get_model_summary(model, dump_input))
if config.TEST.MODEL_FILE:
logger.info('=> loading model from {}'.format(config.TEST.MODEL_FILE))
model.load_state_dict(torch.load(config.TEST.MODEL_FILE))
else:
model_state_file = os.path.join(final_output_dir,
'final_state.pth.tar')
logger.info('=> loading model from {}'.format(model_state_file))
model.load_state_dict(torch.load(model_state_file))
gpus = list(config.GPUS)
model = torch.nn.DataParallel(model, device_ids=gpus).cuda()
# define loss function (criterion) and optimizer
criterion = torch.nn.CrossEntropyLoss().cuda()
# Data loading code
dataset_name = config.DATASET.DATASET
if dataset_name == 'imagenet':
valdir = os.path.join(config.DATASET.ROOT + '/images',
config.DATASET.TEST_SET)
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
transform_valid = transforms.Compose([
transforms.Resize(int(config.MODEL.IMAGE_SIZE[0] / 0.875)),
transforms.CenterCrop(config.MODEL.IMAGE_SIZE[0]),
transforms.ToTensor(),
normalize,
])
valid_dataset = datasets.ImageFolder(valdir, transform_valid)
else:
assert dataset_name == "cifar10", "Only CIFAR-10 is supported at this phase"
normalize = transforms.Normalize((0.4914, 0.4822, 0.4465),
(0.2023, 0.1994, 0.2010))
transform_valid = transforms.Compose([
transforms.ToTensor(),
normalize,
])
valid_dataset = datasets.CIFAR10(root=f'{config.DATASET.ROOT}',
train=False,
download=True,
transform=transform_valid)
valid_loader = torch.utils.data.DataLoader(
valid_dataset,
batch_size=config.TEST.BATCH_SIZE_PER_GPU *
len(gpus) if not check_contract else 1,
shuffle=False,
num_workers=config.WORKERS,
pin_memory=True)
# evaluate on validation set
if not check_contract:
return 'top1', validate(
config,
valid_loader,
model,
criterion,
None,
None,
final_output_dir,
tb_log_dir,
None,
).cpu().numpy()
else:
return 'maxeigen', validate_contractivity(
valid_loader,
model.module.cuda(),
n_iter=n_iter,
).cpu().numpy()