def start_evaluation(args):
device = util.setup_torch()
num_classes = 100 if args.dataset == "cifar100" else 10
train_loader, test_loader = get_cifar(num_classes,
batch_size=args.batch_size)
# for benchmarking, decided whether we want to use unique test folders
if USE_ID:
test_id = util.generate_id()
else:
test_id = ""
results_dir = Path(args.results_dir).joinpath(test_id)
results_dir = Path(results_dir).joinpath(args.dataset)
util.check_dir(results_dir)
# Parsing arguments and prepare settings for training
params = {
"epochs": args.epochs,
"modes": args.modes,
"t_checkpoint": args.t_checkpoint,
"results_dir": results_dir,
"train_loader": train_loader,
"test_loader": test_loader,
"batch_size": args.batch_size,
# model configuration
"device": device,
"teacher_name": args.t_name,
"student_name": args.s_name,
"num_classes": num_classes,
# hyperparameters
"weight_decay": args.weight_decay,
"learning_rate": args.learning_rate,
"momentum": args.momentum,
"sched": "multisteplr",
"optim": "SGD",
# fixed knowledge distillation parameters
"lambda_student": 0.5,
"T_student": 5,
}
test_conf_name = results_dir.joinpath("test_config.json")
util.dump_json_config(test_conf_name, params)
run_benchmarks(args.modes, params, args.s_name, args.t_name)
plot_results(results_dir, test_id=test_id)
device = "cuda" if torch.cuda.is_available() else "cpu"
size = len(dataset.dataset)
model.eval()
correct = 0
with torch.no_grad():
for X, y in dataset:
X, y = X.to(device), y.to(device)
pred = model(X)
correct += (pred.argmax(1) == y).type(torch.float).sum().item()
correct /= size
print(f"Test Error: Accuracy: {(100*correct):>0.1f}%")
return correct
if __name__ == "__main__":
from model_factory import create_cnn_model, is_resnet
from data_loader import get_cifar
from train import load_checkpoint
model_fname = "plain2_Bpodw_best.pth.tar"
name = model_fname.split('_')[0]
model = create_cnn_model(name, 'cifar100', use_cuda=True)
model = load_checkpoint(model, model_fname)
train, test = get_cifar()
evaluate(model, test)
'lambda_student': args.lambda_student,
'milestones': args.milestones,
'gamma': args.gamma
}
# Train Teacher if provided a teacher, otherwise it's a normal training using only cross entropy loss
# This is for training single models(NOKD in paper) for baselines models (or training the first teacher)
if args.teacher:
teacher_model = create_cnn_model(args.teacher, dataset, use_cuda=args.cuda)
if args.teacher_checkpoint:
print("---------- Loading Teacher -------")
teacher_model = load_checkpoint(teacher_model, args.teacher_checkpoint)
else:
print("---------- Training Teacher -------")
train_loader, test_loader = get_cifar(num_classes, args.dataset_dir, args.batch_size, True)
teacher_train_config = copy.deepcopy(train_config)
teacher_name = '{}_{}_best.pth.tar'.format(args.teacher, trial_id)
teacher_train_config['name'] = args.teacher
teacher_trainer = TrainManager(teacher_model, teacher=None, train_loader=train_loader, test_loader=test_loader, train_config=teacher_train_config)
teacher_trainer.train()
teacher_model = load_checkpoint(teacher_model, os.path.join('./', teacher_name))
# Student training
print("---------- Training Student -------")
student_train_config = copy.deepcopy(train_config)
train_loader, test_loader = get_cifar(num_classes, args.dataset_dir, args.batch_size, True)
student_train_config['name'] = args.student
student_trainer = TrainManager(student_model, teacher=teacher_model, train_loader=train_loader, test_loader=test_loader, train_config=student_train_config)
best_student_acc = student_trainer.train()
nni.report_final_result(best_student_acc)
if args.teacher_checkpoint:
print("---------- Loading Teacher -------")
teacher_model = load_checkpoint(teacher_model,
args.teacher_checkpoint)
# train_loader, test_loader = get_cifar(num_classes)
# teacher_train_config = copy.deepcopy(train_config)
# teacher_train_config['Student_name'] = args.teacher
# teacher_train_config['TA_name'] = None
# teacher_trainer = TrainManager(teacher_model, TA=None, teacher=None, train_loader=train_loader,test_loader=test_loader, train_config=teacher_train_config)
# acc = teacher_trainer.validate()
# print('Teacher ACC:', acc)
else:
print("---------- Training Teacher -------")
train_loader, test_loader = get_cifar(num_classes)
teacher_train_config = copy.deepcopy(train_config)
teacher_name = '{}_{}_best.pth.tar'.format(args.teacher, trial_id)
teacher_train_config['Student_name'] = args.teacher
teacher_trainer = TrainManager(teacher_model,
TA=None,
teacher=None,
train_loader=train_loader,
test_loader=test_loader,
train_config=teacher_train_config)
best_teacher_acc, process_form = teacher_trainer.train()
teacher_model = load_checkpoint(teacher_model,
os.path.join('./', teacher_name))
# Student training
print("Teacher param size = %fMB",
'weight_decay': args.weight_decay,
'device': 'cuda:1' if args.cuda else 'cpu',
'is_plane': not is_resnet(args.student),
}
if args.teacher:
teacher_model = create_cnn_model(args.teacher,
dataset,
use_cuda=args.cuda)
if args.teacher_checkpoint:
print('Loading Teacher')
teacher_model = load_checkpoint(teacher_model,
args.teacher_checkpoiint)
else:
print('Training Teacher')
train_loader, test_loader = get_cifar(num_classes)
teacher_train_config = copy.deepcopy(train_config)
teacher_name = '{}_best.pth.tar'.format(args.teacher)
teacher_train_config['name'] = args.teacher
teacher_trainer = TrainManager(teacher_model,
teacher=None,
train_loader=train_loader,
test_loader=test_loader,
train_config=teacher_train_config)
teacher_trainer.train()
teacher_model = load_checkpoint(teacher_model,
os.path.join('./', teacher_name))
print('Trainig Student')
student_train_config = copy.deepcopy(train_config)
train_loader, test_loader = get_cifar(num_classes)
continue
print('student path is: {}'.format(student_path))
# create student model for CIFAR10; usually shake26 for initial student and resnet8 thereafter.
try:
student_model = load_teacher_model(student_path,
'resnet8',
dataset,
cuda_option=args.cuda)
except Exception as e:
print(e)
continue
# where to dump final model
# dataloaders for training paradigm --- do all three for generalization
_, test_loader_cifar = get_cifar(num_classes,
batch_size=args.batch_size,
crop=True)
_, test_loader_imagenet_far = get_imagenet(num_classes,
batch_size=args.batch_size,
crop=True)
_, test_loader_cinic = get_cinic(num_classes,
batch_size=args.batch_size,
crop=True)
print("---------- Evaluating Student -------")
student_eval_cifar = EvalManager(student_model,
test_loader=test_loader_cifar,
eval_config=eval_config)
best_valacc1, best_valloss1 = student_eval_cifar.validate()