device = torch.device('cuda:0')
device_ids = 'cuda'
use_gpu = True
else:
device = torch.device('cpu')
device_ids = 'cpu'
use_gpu = False
if args.data is None:
raise ValueError('you should specify the input data')
if args.out_folder is None:
raise ValueError('you should specify the output folder')
if not os.path.exists(args.out_folder):
os.makedirs(args.out_folder)
model = MLP(in_dim = args.in_dim, hidden_dims = args.hidden_dims, out_dim = args.out_dim, nonlinearity = args.nonlinearity)
criterion = nn.CrossEntropyLoss()
model = model.cuda() if use_gpu else model
criterion = criterion.cuda() if use_gpu else criterion
if args.optim.lower() in ['sgd',]:
optim = torch.optim.SGD(model.parameters(), lr = args.lr, momentum = 0.9, weight_decay = 1e-4)
elif args.optim.lower() in ['adam',]:
optim = torch.optim.Adam(model.parameters(), lr = args.lr, betas = (0.9, 0.99), weight_decay = 1e-4)
else:
raise ValueError('Unrecognized Optimizer: %s' % args.optim.lower())
data_loader = load_pkl(args.data, batch_size = args.batch_size)
setup_config = {kwarg: value for kwarg, value in args._get_kwargs()}
for idx in range(args.total_iters):
device_ids = 'cuda'
use_gpu = True
else:
device = torch.device('cpu')
device_ids = 'cpu'
use_gpu = False
if args.data is None:
raise ValueError('you should specify the input data')
if args.out_folder is None:
raise ValueError('you should specify the output folder')
if not os.path.exists(args.out_folder):
os.makedirs(args.out_folder)
model = MLP(in_dim=args.in_dim,
hidden_dims=args.hidden_dims,
out_dim=args.out_dim,
nonlinearity=args.nonlinearity)
criterion = nn.CrossEntropyLoss()
model = model.cuda() if use_gpu else model
criterion = criterion.cuda() if use_gpu else criterion
if args.optim.lower() in [
'sgd',
]:
optimizer = torch.optim.SGD(model.parameters(),
lr=args.lr,
momentum=0.9,
weight_decay=1e-4)
elif args.optim.lower() in [
'adam',
]:
elif args.dataset.lower() in [
'cifar10',
]:
train_loader, test_loader = cifar10(batch_size=args.batch_size,
batch_size_test=args.batch_size,
normalization=args.normalization)
else:
raise ValueError('Unrecognized dataset: %s' % args.dataset)
# Parse IO
if not os.path.exists(args.out_folder):
os.makedirs(args.out_folder)
# Parse model
model = MLP(in_dim=args.in_dim,
hidden_dims=args.hidden_dims,
out_dim=args.out_dim,
nonlinearity=args.nonlinearity)
criterion = nn.CrossEntropyLoss()
model = model.cuda(device) if use_gpu else model
criterion = criterion.cuda(device) if use_gpu else criterion
# load pretrained model
if args.model2load != None:
if args.model2load.endswith('.ckpt'):
model = load_ckpt_model(model=model,
name='fc1',
ckpt_file=args.model2load,
device=device,
param2load=None)
elif args.model2load.endswith('.mat'):
model = load_mat_model(model=model,
], 'Subset tag can be only "train" or "test", but %s found' % args.subset
data_loader = train_loader if args.subset in [
'train',
] else test_loader
# Parse IO
assert os.path.exists(
args.model2load
), 'model2load file %s does not exists' % args.model2load
out_dir = os.path.dirname(args.out_file)
if out_dir != '' and not os.path.exists(out_dir):
os.makedirs(out_dir)
# Parse model
model = MLP(in_dim=args.in_dim,
hidden_dims=args.hidden_dims,
out_dim=args.out_dim,
nonlinearity=args.nonlinearity)
model = model.cuda(device) if use_gpu else model
if args.model2load.endswith('.ckpt'):
ckpt = torch.load(args.model2load)
model.load_state_dict(ckpt)
elif args.model2load.endswith('.mat'):
model = load_mat_model(model=model,
name='fc1',
mat_file=args.model2load,
device=device)
elif args.model2load.endswith('.pth'):
model = load_pth_model(model=model,
name='fc1',
pth_file=args.model2load,
device=device)
]:
data_loader = load_fmnist(batch_size=args.batch_size,
dset='test',
subset=args.subset)
elif args.dataset.lower() in [
'svhn',
]:
data_loader = load_svhn(batch_size=args.batch_size,
dset='test',
subset=args.subset)
else:
raise ValueError('Unrecognized dataset: %s' % args.dataset.lower())
# Model configuration
model = MLP(in_dim=args.in_dim,
hidden_dims=args.hidden_dims,
out_dim=args.out_dim,
nonlinearity=args.nonlinearity)
model = model.cuda(device) if use_gpu else model
ckpt = torch.load(args.model2load)
model.load_state_dict(ckpt)
model.eval()
# Configure the certification parameter
init_value = np.sqrt(args.init_margin)
norm = np.inf if args.norm <= 0 else args.norm
var_list = var_init(mode=args.mode,
batch_size=args.batch_size,
in_dim=args.in_dim,
init_value=init_value,
device=device)