if args.dataset == 'cifar10':
train_loader, test_loader = dataset.get_cifar10(batch_size=args.batch_size,
num_workers=1)
elif args.dataset == 'cifar100':
train_loader, test_loader = dataset.get_cifar100(
batch_size=args.batch_size, num_workers=1)
elif args.dataset == 'imagenet':
train_loader, test_loader = dataset.get_imagenet(
batch_size=args.batch_size, num_workers=1)
else:
raise ValueError("Unknown dataset type")
assert args.model in ['VGG8', 'DenseNet40', 'ResNet18'], args.model
if args.model == 'VGG8':
from models import VGG
model = VGG.vgg8(args=args, logger=logger)
criterion = wage_util.SSE()
elif args.model == 'DenseNet40':
from models import DenseNet
model = DenseNet.densenet40(args=args, logger=logger)
criterion = wage_util.SSE()
elif args.model == 'ResNet18':
from models import ResNet
model = ResNet.resnet18(args=args, logger=logger)
criterion = torch.nn.CrossEntropyLoss()
else:
raise ValueError("Unknown model type")
if args.cuda:
model.cuda()
# data loader and model
assert args.dataset in ['cifar10', 'cifar100', 'imagenet'], args.dataset
if args.dataset == 'cifar10':
train_loader, test_loader = dataset.get_cifar10(batch_size=args.batch_size, num_workers=1)
elif args.dataset == 'cifar100':
train_loader, test_loader = dataset.get_cifar100(batch_size=args.batch_size, num_workers=1)
elif args.dataset == 'imagenet':
train_loader, test_loader = dataset.get_imagenet(batch_size=args.batch_size, num_workers=1)
else:
raise ValueError("Unknown dataset type")
assert args.model in ['VGG8', 'DenseNet40', 'ResNet18'], args.model
if args.model == 'VGG8':
from models import VGG
model_path = './log/VGG8.pth' # WAGE mode pretrained model
modelCF = VGG.vgg8(args = args, logger=logger, pretrained = model_path)
elif args.model == 'DenseNet40':
from models import DenseNet
model_path = './log/DenseNet40.pth' # WAGE mode pretrained model
modelCF = DenseNet.densenet40(args = args, logger=logger, pretrained = model_path)
elif args.model == 'ResNet18':
from models import ResNet
# FP mode pretrained model, loaded from 'https://download.pytorch.org/models/resnet18-5c106cde.pth'
# model_path = './log/xxx.pth'
# modelCF = ResNet.resnet18(args = args, logger=logger, pretrained = model_path)
modelCF = ResNet.resnet18(args = args, logger=logger, pretrained = True)
else:
raise ValueError("Unknown model type")
if args.cuda:
modelCF.cuda()
