if args.dataset.lower()=='cifar100':
num_classes=100
elif args.dataset.lower()=='imagenet':
num_classes=1000
elif args.dataset.lower()=='tinyimagenet':
num_classes=200
else:
num_classes=10
net, model_name, Q = instantiate_model(dataset=args.dataset,
num_classes=num_classes,
input_quant=args.input_quant,
arch=args.arch,
dorefa=args.dorefa,
abit=args.abit,
wbit=args.wbit,
qin=args.qin,
qout=args.qout,
suffix=args.suffix,
load=args.pretrained,
torch_weights=args.torch_weights,
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu'))
framework = Framework(net=net,
adversarial_testing=args.adv_tst,
model_name=model_name,
preprocess=Q,
normalize=args.normalize,
dataset=args.dataset,
train_batch_size=args.train_batch_size,
if args.dataset.lower()=='imagenet':
num_classes=1000
elif args.dataset.lower()=='tinyimagenet':
num_classes=200
elif args.dataset.lower()=='cifar100':
num_classes=100
else:
num_classes=10
net, model_name, Q = instantiate_model(dataset=args.dataset,
num_classes=num_classes,
input_quant=args.input_quant,
arch=args.arch,
dorefa=args.dorefa,
abit=args.abit,
wbit=args.wbit,
qin=args.qin,
qout=args.qout,
suffix=args.suffix,
load=args.pretrained,
torch_weights=args.torch_weights,
normalize= args.normalize)
# default `log_dir` is "runs" - we'll be more specific here
writer = SummaryWriter('./pretrained/'+args.dataset.lower()+'/runs/'+model_name)
framework = Framework(net=net,
model_name=model_name,
preprocess=Q,
dataset=args.dataset,
epochs=args.epochs,
num_classes = 1000
elif args.dataset.lower() == 'tinyimagenet':
num_classes = 200
elif args.dataset.lower() == 'cifar100':
num_classes = 100
else:
num_classes = 10
net, model_name, Q = instantiate_model(
dataset=args.dataset,
num_classes=num_classes,
input_quant=args.input_quant,
arch=args.arch,
dorefa=args.dorefa,
abit=args.abit,
wbit=args.wbit,
qin=args.qin,
qout=args.qout,
suffix=args.suffix,
load=args.pretrained,
torch_weights=args.torch_weights,
device=torch.device('cuda' if torch.cuda.is_available() else 'cpu'),
normalize=args.normalize)
framework = Framework(net=net,
model_name=model_name,
preprocess=Q,
dataset=args.dataset,
train_batch_size=args.train_batch_size,
test_batch_size=args.test_batch_size,
val_split=args.val_split,