# set number of classes
if args.dataset == 'CIFAR10':
num_classes = 10
model_folder = "model_cifar"
root='/home/chendefang/MC/Data'
elif args.dataset == 'CIFAR100':
num_classes = 100
model_folder = "model_cifar"
root='/home/chendefang/MC/Data'
elif args.dataset == 'imagenet':
num_classes = 1000
model_folder = "model_imagenet"
root = '/home/meijianping/Test/Data'
# Load data
train_loader, test_loader = data_loader.dataloader(data_name = args.dataset, batch_size = args.batch_size, num_workers = args.num_workers, root=root)
logging.info("- Done.")
# Training from scratch
model_fd = getattr(models, model_folder)
if args.MulStu:
model_cfg = getattr(model_fd, 'MultiNet')
model = getattr(model_cfg, 'StuNet')(model = args.model, num_branches = args.num_branches, num_classes = num_classes, input_channel=utils.lookup(args.model), dropout = args.dropout)
elif args.type == 'DML':
model_cfg = getattr(model_fd, 'DML')
model = getattr(model_cfg, 'MutualNet')(model = args.model, num_branches = args.num_branches, num_classes = num_classes)
else:
if "resnet" in args.model:
model_cfg = getattr(model_fd, 'resnet_GL')
model = getattr(model_cfg, args.model)(num_classes = num_classes, num_branches = args.num_branches, input_channel=utils.lookup(args.model))
elif "vgg" in args.model:
# set number of classes
if args.dataset == 'CIFAR10':
num_classes = 10
model_folder = "model_cifar"
root = './Data'
elif args.dataset == 'CIFAR100':
num_classes = 100
model_folder = "model_cifar"
root = './Data'
elif args.dataset == 'imagenet':
num_classes = 1000
model_folder = "model_imagenet"
root = './Data'
# Load data
train_loader, test_loader = data_loader.dataloader(
data_name=args.dataset, batch_size=args.batch_size, root=root)
logging.info("- Done.")
# Training from scratch for student model
model_fd = getattr(models, model_folder)
if "resnet" in args.model:
model_cfg = getattr(model_fd, 'resnet')
model = getattr(model_cfg, args.model)(num_classes=num_classes)
elif "vgg" in args.model:
model_cfg = getattr(model_fd, 'vgg')
model = getattr(model_cfg, args.model)(num_classes=num_classes,
dropout=args.dropout)
elif "densenet" in args.model:
model_cfg = getattr(model_fd, 'densenet')
model = getattr(model_cfg, args.model)(num_classes=num_classes)