def main():
best_acc = 0
opt = parse_option()
# tensorboard logger
logger = tb_logger.Logger(logdir=opt.tb_folder, flush_secs=2)
# dataloader
if opt.dataset == 'cifar100':
if opt.distill in ['crd', 'ml_cpc', 'cpc']:
train_loader, val_loader, n_data = get_cifar100_dataloaders_sample(
batch_size=opt.batch_size,
num_workers=opt.num_workers,
k=opt.nce_k,
mode=opt.mode)
else:
train_loader, val_loader, n_data = get_cifar100_dataloaders(
batch_size=opt.batch_size,
num_workers=opt.num_workers,
is_instance=True)
n_cls = 100
else:
raise NotImplementedError(opt.dataset)
# model
model_t = load_teacher(opt.path_t, n_cls)
model_s = model_dict[opt.model_s](num_classes=n_cls)
data = torch.randn(2, 3, 32, 32)
model_t.eval()
model_s.eval()
feat_t, _ = model_t(data, is_feat=True)
feat_s, _ = model_s(data, is_feat=True)
module_list = nn.ModuleList([])
module_list.append(model_s)
trainable_list = nn.ModuleList([])
trainable_list.append(model_s)
criterion_cls = nn.CrossEntropyLoss()
criterion_div = DistillKL(opt.kd_T)
if opt.distill == 'kd':
criterion_kd = DistillKL(opt.kd_T)
elif opt.distill == 'hint':
criterion_kd = HintLoss()
regress_s = ConvReg(feat_s[opt.hint_layer].shape,
feat_t[opt.hint_layer].shape)
module_list.append(regress_s)
trainable_list.append(regress_s)
elif opt.distill == 'crd':
opt.s_dim = feat_s[-1].shape[1]
opt.t_dim = feat_t[-1].shape[1]
opt.n_data = n_data
criterion_kd = CRDLoss(opt)
module_list.append(criterion_kd.embed_s)
module_list.append(criterion_kd.embed_t)
trainable_list.append(criterion_kd.embed_s)
trainable_list.append(criterion_kd.embed_t)
elif opt.distill == 'cpc':
opt.s_dim = feat_s[-1].shape[1]
opt.t_dim = feat_t[-1].shape[1]
opt.n_data = n_data
criterion_kd = CPCLoss(opt)
module_list.append(criterion_kd.embed_s)
module_list.append(criterion_kd.embed_t)
trainable_list.append(criterion_kd.embed_s)
trainable_list.append(criterion_kd.embed_t)
elif opt.distill == 'ml_cpc':
opt.s_dim = feat_s[-1].shape[1]
opt.t_dim = feat_t[-1].shape[1]
opt.n_data = n_data
criterion_kd = MLCPCLoss(opt)
module_list.append(criterion_kd.embed_s)
module_list.append(criterion_kd.embed_t)
trainable_list.append(criterion_kd.embed_s)
trainable_list.append(criterion_kd.embed_t)
elif opt.distill == 'attention':
criterion_kd = Attention()
elif opt.distill == 'nst':
criterion_kd = NSTLoss()
elif opt.distill == 'similarity':
criterion_kd = Similarity()
elif opt.distill == 'rkd':
criterion_kd = RKDLoss()
elif opt.distill == 'pkt':
criterion_kd = PKT()
elif opt.distill == 'kdsvd':
criterion_kd = KDSVD()
elif opt.distill == 'correlation':
criterion_kd = Correlation()
embed_s = LinearEmbed(feat_s[-1].shape[1], opt.feat_dim)
embed_t = LinearEmbed(feat_t[-1].shape[1], opt.feat_dim)
module_list.append(embed_s)
module_list.append(embed_t)
trainable_list.append(embed_s)
trainable_list.append(embed_t)
elif opt.distill == 'vid':
s_n = [f.shape[1] for f in feat_s[1:-1]]
t_n = [f.shape[1] for f in feat_t[1:-1]]
criterion_kd = nn.ModuleList(
[VIDLoss(s, t, t) for s, t in zip(s_n, t_n)])
# add this as some parameters in VIDLoss need to be updated
trainable_list.append(criterion_kd)
elif opt.distill == 'abound':
s_shapes = [f.shape for f in feat_s[1:-1]]
t_shapes = [f.shape for f in feat_t[1:-1]]
connector = Connector(s_shapes, t_shapes)
# init stage training
init_trainable_list = nn.ModuleList([])
init_trainable_list.append(connector)
init_trainable_list.append(model_s.get_feat_modules())
criterion_kd = ABLoss(len(feat_s[1:-1]))
init(model_s, model_t, init_trainable_list, criterion_kd, train_loader,
logger, opt)
# classification
module_list.append(connector)
elif opt.distill == 'factor':
s_shape = feat_s[-2].shape
t_shape = feat_t[-2].shape
paraphraser = Paraphraser(t_shape)
translator = Translator(s_shape, t_shape)
# init stage training
init_trainable_list = nn.ModuleList([])
init_trainable_list.append(paraphraser)
criterion_init = nn.MSELoss()
init(model_s, model_t, init_trainable_list, criterion_init,
train_loader, logger, opt)
# classification
criterion_kd = FactorTransfer()
module_list.append(translator)
module_list.append(paraphraser)
trainable_list.append(translator)
elif opt.distill == 'fsp':
s_shapes = [s.shape for s in feat_s[:-1]]
t_shapes = [t.shape for t in feat_t[:-1]]
criterion_kd = FSP(s_shapes, t_shapes)
# init stage training
init_trainable_list = nn.ModuleList([])
init_trainable_list.append(model_s.get_feat_modules())
init(model_s, model_t, init_trainable_list, criterion_kd, train_loader,
logger, opt)
# classification training
pass
else:
raise NotImplementedError(opt.distill)
criterion_list = nn.ModuleList([])
criterion_list.append(criterion_cls) # classification loss
criterion_list.append(
criterion_div) # KL divergence loss, original knowledge distillation
criterion_list.append(criterion_kd) # other knowledge distillation loss
# optimizer
optimizer = optim.SGD(trainable_list.parameters(),
lr=opt.learning_rate,
momentum=opt.momentum,
weight_decay=opt.weight_decay)
# append teacher after optimizer to avoid weight_decay
module_list.append(model_t)
if torch.cuda.is_available():
module_list.cuda()
criterion_list.cuda()
cudnn.benchmark = True
# validate teacher accuracy
teacher_acc, _, _ = validate(val_loader, model_t, criterion_cls, opt)
print('teacher accuracy: ', teacher_acc)
# routine
for epoch in range(1, opt.epochs + 1):
adjust_learning_rate(epoch, opt, optimizer)
print("==> training...")
time1 = time.time()
train_acc, train_loss = train(epoch, train_loader, module_list,
criterion_list, optimizer, opt)
time2 = time.time()
print('epoch {}, total time {:.2f}'.format(epoch, time2 - time1))
logger.log_value('train_acc', train_acc, epoch)
logger.log_value('train_loss', train_loss, epoch)
test_acc, tect_acc_top5, test_loss = validate(val_loader, model_s,
criterion_cls, opt)
logger.log_value('test_acc', test_acc, epoch)
logger.log_value('test_loss', test_loss, epoch)
logger.log_value('test_acc_top5', tect_acc_top5, epoch)
# save the best model
if test_acc > best_acc:
best_acc = test_acc
state = {
'epoch': epoch,
'model': model_s.state_dict(),
'best_acc': best_acc,
}
save_file = os.path.join(opt.save_folder,
'{}_best.pth'.format(opt.model_s))
print('saving the best model!')
torch.save(state, save_file)
# regular saving
if epoch % opt.save_freq == 0:
print('==> Saving...')
state = {
'epoch': epoch,
'model': model_s.state_dict(),
'accuracy': test_acc,
}
save_file = os.path.join(
opt.save_folder, 'ckpt_epoch_{epoch}.pth'.format(epoch=epoch))
torch.save(state, save_file)
# This best accuracy is only for printing purpose.
# The results reported in the paper/README is from the last epoch.
print('best accuracy:', best_acc)
# save model
state = {
'opt': opt,
'model': model_s.state_dict(),
}
save_file = os.path.join(opt.save_folder,
'{}_last.pth'.format(opt.model_s))
torch.save(state, save_file)
def main():
best_acc = 0
opt = parse_option()
# tensorboard logger
logger = tb_logger.Logger(logdir=opt.tb_folder, flush_secs=2)
if os.path.exists('./save/log/'):
pass
else:
os.mkdir('./save/log/')
log = open(os.path.join('./save/log/', 'log_{}.txt'.format(opt.path_config[-11:-7])), 'w')
print_log('save path : {}'.format("./save/"), log)
# dataloader
if opt.dataset == 'imagenet':
if opt.distill in ['crd']:
train_loader, val_loader, n_data = get_dataloader_sample(batch_size=opt.batch_size,
num_workers=opt.num_workers,
k=opt.nce_k,
mode=opt.mode)
else:
train_loader, val_loader, n_data = get_imagenet_dataloader(batch_size=opt.batch_size,
num_workers=opt.num_workers,
is_instance=True)
'''import torchvision.datasets as dset
import torchvision.transforms as transforms
data_folder = '/gdata/ImageNet2012'
traindir = os.path.join(data_folder, 'train')
valdir = os.path.join(data_folder, 'val')
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
train_data1 = dset.ImageFolder(traindir, transforms.Compose([
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize,
]))
valid_data = dset.ImageFolder(valdir, transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
normalize,
]))
train_loader = torch.utils.data.DataLoader(
train_data1, batch_size=opt.batch_size, shuffle=True,
pin_memory=True, num_workers=opt.workers)
val_loader = torch.utils.data.DataLoader(
valid_data, batch_size=opt.batch_size, shuffle=True,
pin_memory=True, num_workers=opt.workers)
n_data = len(train_data1)'''
n_cls = 1000
else:
raise NotImplementedError(opt.dataset)
print("##")
# model
model_t = EfficientNet.from_pretrained('efficientnet-b0',weights_path='./pretrain_efficientNet/pretrain_efficientNet.pth')
#torch.save(model_t.state_dict(), './pretrain_efficientNet/pretrain_efficientNet.pt')
from proxyless_nas.jj import get_proxyless_model
model_s = get_proxyless_model(net_config_path=opt.path_config)
gpus = [0, 1, 2, 3]
torch.cuda.set_device('cuda:{}'.format(gpus[0]))
data = torch.randn(2, 3, 224, 224)
model_t.eval()
model_s.eval()
feat_t, _ = model_t(data, is_feat=True)
feat_s, _ = model_s(data, is_feat=True)
module_list = nn.ModuleList([])
module_list.append(model_s)
trainable_list = nn.ModuleList([])
trainable_list.append(model_s)
criterion_cls = nn.CrossEntropyLoss()
criterion_div = DistillKL(opt.kd_T)
if opt.distill == 'kd':
criterion_kd = DistillKL(opt.kd_T)
elif opt.distill == 'hint':
criterion_kd = HintLoss()
regress_s = ConvReg(feat_s[opt.hint_layer].shape, feat_t[opt.hint_layer].shape)
module_list.append(regress_s)
trainable_list.append(regress_s)
elif opt.distill == 'crd':
opt.s_dim = feat_s[-1].shape[1]
opt.t_dim = feat_t[-1].shape[1]
opt.n_data = n_data
criterion_kd = CRDLoss(opt)
module_list.append(criterion_kd.embed_s)
module_list.append(criterion_kd.embed_t)
trainable_list.append(criterion_kd.embed_s)
trainable_list.append(criterion_kd.embed_t)
elif opt.distill == 'attention':
criterion_kd = Attention()
elif opt.distill == 'nst':
criterion_kd = NSTLoss()
elif opt.distill == 'similarity':
criterion_kd = Similarity()
elif opt.distill == 'rkd':
criterion_kd = RKDLoss()
elif opt.distill == 'pkt':
criterion_kd = PKT()
elif opt.distill == 'kdsvd':
criterion_kd = KDSVD()
elif opt.distill == 'correlation':
criterion_kd = Correlation()
embed_s = LinearEmbed(feat_s[-1].shape[1], opt.feat_dim)
embed_t = LinearEmbed(feat_t[-1].shape[1], opt.feat_dim)
module_list.append(embed_s)
module_list.append(embed_t)
trainable_list.append(embed_s)
trainable_list.append(embed_t)
elif opt.distill == 'vid':
s_n = [f.shape[1] for f in feat_s[1:-1]]
t_n = [f.shape[1] for f in feat_t[1:-1]]
criterion_kd = nn.ModuleList(
[VIDLoss(s, t, t) for s, t in zip(s_n, t_n)]
)
# add this as some parameters in VIDLoss need to be updated
trainable_list.append(criterion_kd)
elif opt.distill == 'abound':
s_shapes = [f.shape for f in feat_s[1:-1]]
t_shapes = [f.shape for f in feat_t[1:-1]]
connector = Connector(s_shapes, t_shapes)
# init stage training
init_trainable_list = nn.ModuleList([])
init_trainable_list.append(connector)
init_trainable_list.append(model_s.get_feat_modules())
criterion_kd = ABLoss(len(feat_s[1:-1]))
init(model_s, model_t, init_trainable_list, criterion_kd, train_loader, logger, opt)
# classification
module_list.append(connector)
elif opt.distill == 'factor':
s_shape = feat_s[-2].shape
t_shape = feat_t[-2].shape
paraphraser = Paraphraser(t_shape)
translator = Translator(s_shape, t_shape)
# init stage training
init_trainable_list = nn.ModuleList([])
init_trainable_list.append(paraphraser)
criterion_init = nn.MSELoss()
init(model_s, model_t, init_trainable_list, criterion_init, train_loader, logger, opt)
# classification
criterion_kd = FactorTransfer()
module_list.append(translator)
module_list.append(paraphraser)
trainable_list.append(translator)
elif opt.distill == 'fsp':
s_shapes = [s.shape for s in feat_s[:-1]]
t_shapes = [t.shape for t in feat_t[:-1]]
criterion_kd = FSP(s_shapes, t_shapes)
# init stage training
init_trainable_list = nn.ModuleList([])
init_trainable_list.append(model_s.get_feat_modules())
init(model_s, model_t, init_trainable_list, criterion_kd, train_loader, logger, opt)
# classification training
pass
else:
raise NotImplementedError(opt.distill)
criterion_list = nn.ModuleList([])
criterion_list.append(criterion_cls) # classification loss
criterion_list.append(criterion_div) # KL divergence loss, original knowledge distillation
criterion_list.append(criterion_kd) # other knowledge distillation loss
# optimizer
optimizer = optim.SGD(trainable_list.parameters(),
lr=opt.learning_rate,
momentum=opt.momentum,
weight_decay=opt.weight_decay)
# append teacher after optimizer to avoid weight_decay
module_list.append(model_t)
if torch.cuda.is_available():
module_list.cuda()
criterion_list.cuda()
cudnn.benchmark = True
# validate teacher accuracy
teacher_acc, _, _ = validate(val_loader, model_t, criterion_cls, opt,log)
print('teacher accuracy: ', teacher_acc)
print_log("teacher accuracy:{}".format(teacher_acc), log)
# routine
for epoch in range(1, opt.epochs + 1):
adjust_learning_rate(epoch, opt, optimizer)
print("==> training...")
print_log("==> training...", log)
time1 = time.time()
train_acc, train_loss = train(epoch, train_loader, module_list, criterion_list, optimizer, opt,log)
time2 = time.time()
print('epoch {}, total time {:.2f}'.format(epoch, time2 - time1))
print_log('epoch {}, total time {:.2f}'.format(epoch, time2 - time1), log)
logger.log_value('train_acc', train_acc, epoch)
logger.log_value('train_loss', train_loss, epoch)
test_acc, tect_acc_top5, test_loss = validate(val_loader, model_s, criterion_cls, opt,log)
logger.log_value('test_acc', test_acc, epoch)
logger.log_value('test_loss', test_loss, epoch)
logger.log_value('test_acc_top5', tect_acc_top5, epoch)
# save the best model
if test_acc > best_acc:
best_acc = test_acc
state = {
'epoch': epoch,
'model': model_s.state_dict(),
'best_acc': best_acc,
}
save_file = os.path.join(opt.save_folder, '{}_best.pth'.format(opt.model_s))
print('saving the best model!')
print_log('saving the best model!', log)
torch.save(state, save_file)
# regular saving
if epoch % opt.save_freq == 0 or epoch<10:
print('==> Saving...')
state = {
'epoch': epoch,
'model': model_s.state_dict(),
'accuracy': test_acc,
}
save_file = os.path.join(opt.save_folder, 'ckpt_epoch_{epoch}.pth'.format(epoch=epoch))
torch.save(state, save_file)
# This best accuracy is only for printing purpose.
# The results reported in the paper/README is from the last epoch.
print('best accuracy:', best_acc)
print_log('best accuracy:{}'.format(best_acc), log)
# save model
state = {
'opt': opt,
'model': model_s.state_dict(),
}
save_file = os.path.join(opt.save_folder, '{}_last.pth'.format(opt.model_s))
torch.save(state, save_file)
log.close()
def main():
best_acc = 0
opt = parse_option()
torch.manual_seed(2021)
torch.cuda.manual_seed(2021)
torch.backends.cudnn.deterministic = True
# dataloader
if opt.distill in ['crd']:
train_loader, val_loader, n_data = get_cifar100_dataloaders_sample(
opt.data_path,
batch_size=opt.batch_size,
num_workers=opt.num_workers,
k=opt.nce_k,
mode=opt.mode,
use_fake_data=opt.use_fake_data,
fake_data_folder=opt.fake_data_path,
nfake=opt.nfake)
else:
train_loader, val_loader, n_data = get_cifar100_dataloaders(
opt.data_path,
batch_size=opt.batch_size,
num_workers=opt.num_workers,
is_instance=True,
use_fake_data=opt.use_fake_data,
fake_data_folder=opt.fake_data_path,
nfake=opt.nfake)
n_cls = 100
# model
model_t = load_teacher(opt.path_t, n_cls)
model_s = model_dict[opt.model_s](num_classes=n_cls)
## student model name, how to initialize student model, etc.
student_model_filename = 'S_{}_T_{}_{}_r_{}_a_{}_b_{}_epoch_{}'.format(
opt.model_s, opt.model_t, opt.distill, opt.gamma, opt.alpha, opt.beta,
opt.epochs)
if opt.finetune:
ckpt_cnn_filename = os.path.join(
opt.save_folder,
student_model_filename + '_finetune_True_last.pth')
## load pre-trained model
checkpoint = torch.load(opt.init_student_path)
model_s.load_state_dict(checkpoint['model'])
else:
ckpt_cnn_filename = os.path.join(opt.save_folder,
student_model_filename + '_last.pth')
print('\n ' + ckpt_cnn_filename)
data = torch.randn(2, 3, 32, 32)
model_t.eval()
model_s.eval()
feat_t, _ = model_t(data, is_feat=True)
feat_s, _ = model_s(data, is_feat=True)
module_list = nn.ModuleList([])
module_list.append(model_s)
trainable_list = nn.ModuleList([])
trainable_list.append(model_s)
criterion_cls = nn.CrossEntropyLoss()
criterion_div = DistillKL(opt.kd_T)
if opt.distill == 'kd':
criterion_kd = DistillKL(opt.kd_T)
elif opt.distill == 'hint':
criterion_kd = HintLoss()
regress_s = ConvReg(feat_s[opt.hint_layer].shape,
feat_t[opt.hint_layer].shape)
module_list.append(regress_s)
trainable_list.append(regress_s)
elif opt.distill == 'crd':
opt.s_dim = feat_s[-1].shape[1]
opt.t_dim = feat_t[-1].shape[1]
opt.n_data = n_data
criterion_kd = CRDLoss(opt)
module_list.append(criterion_kd.embed_s)
module_list.append(criterion_kd.embed_t)
trainable_list.append(criterion_kd.embed_s)
trainable_list.append(criterion_kd.embed_t)
elif opt.distill == 'attention':
criterion_kd = Attention()
elif opt.distill == 'nst':
criterion_kd = NSTLoss()
elif opt.distill == 'similarity':
criterion_kd = Similarity()
elif opt.distill == 'rkd':
criterion_kd = RKDLoss()
elif opt.distill == 'pkt':
criterion_kd = PKT()
elif opt.distill == 'kdsvd':
criterion_kd = KDSVD()
elif opt.distill == 'correlation':
criterion_kd = Correlation()
embed_s = LinearEmbed(feat_s[-1].shape[1], opt.feat_dim)
embed_t = LinearEmbed(feat_t[-1].shape[1], opt.feat_dim)
module_list.append(embed_s)
module_list.append(embed_t)
trainable_list.append(embed_s)
trainable_list.append(embed_t)
elif opt.distill == 'vid':
s_n = [f.shape[1] for f in feat_s[1:-1]]
t_n = [f.shape[1] for f in feat_t[1:-1]]
criterion_kd = nn.ModuleList(
[VIDLoss(s, t, t) for s, t in zip(s_n, t_n)])
# add this as some parameters in VIDLoss need to be updated
trainable_list.append(criterion_kd)
elif opt.distill == 'abound':
s_shapes = [f.shape for f in feat_s[1:-1]]
t_shapes = [f.shape for f in feat_t[1:-1]]
connector = Connector(s_shapes, t_shapes)
# init stage training
init_trainable_list = nn.ModuleList([])
init_trainable_list.append(connector)
init_trainable_list.append(model_s.get_feat_modules())
criterion_kd = ABLoss(len(feat_s[1:-1]))
init(model_s, model_t, init_trainable_list, criterion_kd, train_loader,
opt)
# classification
module_list.append(connector)
elif opt.distill == 'factor':
s_shape = feat_s[-2].shape
t_shape = feat_t[-2].shape
paraphraser = Paraphraser(t_shape)
translator = Translator(s_shape, t_shape)
# init stage training
init_trainable_list = nn.ModuleList([])
init_trainable_list.append(paraphraser)
criterion_init = nn.MSELoss()
init(model_s, model_t, init_trainable_list, criterion_init,
train_loader, opt)
# classification
criterion_kd = FactorTransfer()
module_list.append(translator)
module_list.append(paraphraser)
trainable_list.append(translator)
elif opt.distill == 'fsp':
s_shapes = [s.shape for s in feat_s[:-1]]
t_shapes = [t.shape for t in feat_t[:-1]]
criterion_kd = FSP(s_shapes, t_shapes)
# init stage training
init_trainable_list = nn.ModuleList([])
init_trainable_list.append(model_s.get_feat_modules())
init(model_s, model_t, init_trainable_list, criterion_kd, train_loader,
opt)
# classification training
pass
else:
raise NotImplementedError(opt.distill)
criterion_list = nn.ModuleList([])
criterion_list.append(criterion_cls) # classification loss
criterion_list.append(
criterion_div) # KL divergence loss, original knowledge distillation
criterion_list.append(criterion_kd) # other knowledge distillation loss
# optimizer
optimizer = optim.SGD(trainable_list.parameters(),
lr=opt.learning_rate,
momentum=opt.momentum,
weight_decay=opt.weight_decay)
# append teacher after optimizer to avoid weight_decay
module_list.append(model_t)
if torch.cuda.is_available():
module_list.cuda()
criterion_list.cuda()
cudnn.benchmark = True
# validate teacher accuracy
teacher_acc, _, _ = validate(val_loader, model_t, criterion_cls, opt)
print('teacher accuracy: ', teacher_acc)
if not os.path.isfile(ckpt_cnn_filename):
print("\n Start training the {} >>>".format(opt.model_s))
## resume training
if opt.resume_epoch > 0:
save_file = opt.save_intrain_folder + "/ckpt_{}_epoch_{}.pth".format(
opt.model_s, opt.resume_epoch)
checkpoint = torch.load(save_file)
model_s.load_state_dict(checkpoint['model'])
optimizer.load_state_dict(checkpoint['optimizer'])
module_list.load_state_dict(checkpoint['module_list'])
trainable_list.load_state_dict(checkpoint['trainable_list'])
criterion_list.load_state_dict(checkpoint['criterion_list'])
# module_list = checkpoint['module_list']
# criterion_list = checkpoint['criterion_list']
# # trainable_list = checkpoint['trainable_list']
# ckpt_test_accuracy = checkpoint['accuracy']
# ckpt_epoch = checkpoint['epoch']
# print('\n Resume training: epoch {}, test_acc {}...'.format(ckpt_epoch, ckpt_test_accuracy))
if torch.cuda.is_available():
module_list.cuda()
criterion_list.cuda()
#end if
for epoch in range(opt.resume_epoch, opt.epochs):
adjust_learning_rate(epoch, opt, optimizer)
print("==> training...")
time1 = time.time()
train_acc, train_loss = train(epoch, train_loader, module_list,
criterion_list, optimizer, opt)
time2 = time.time()
print('epoch {}, total time {:.2f}'.format(epoch, time2 - time1))
test_acc, tect_acc_top5, test_loss = validate(
val_loader, model_s, criterion_cls, opt)
# regular saving
if (epoch + 1) % opt.save_freq == 0:
print('==> Saving...')
state = {
'epoch': epoch,
'model': model_s.state_dict(),
'optimizer': optimizer.state_dict(),
'module_list': module_list.state_dict(),
'criterion_list': criterion_list.state_dict(),
'trainable_list': trainable_list.state_dict(),
'accuracy': test_acc,
}
save_file = os.path.join(
opt.save_intrain_folder,
'ckpt_{}_epoch_{}.pth'.format(opt.model_s, epoch + 1))
torch.save(state, save_file)
##end for epoch
# store model
torch.save({
'opt': opt,
'model': model_s.state_dict(),
}, ckpt_cnn_filename)
print("\n End training CNN.")
else:
print("\n Loading pre-trained {}.".format(opt.model_s))
checkpoint = torch.load(ckpt_cnn_filename)
model_s.load_state_dict(checkpoint['model'])
test_acc, test_acc_top5, _ = validate(val_loader, model_s, criterion_cls,
opt)
print("\n {}, test_acc:{:.3f}, test_acc_top5:{:.3f}.".format(
opt.model_s, test_acc, test_acc_top5))
eval_results_fullpath = opt.save_folder + "/test_result_" + opt.model_name + ".txt"
if not os.path.isfile(eval_results_fullpath):
eval_results_logging_file = open(eval_results_fullpath, "w")
eval_results_logging_file.close()
with open(eval_results_fullpath, 'a') as eval_results_logging_file:
eval_results_logging_file.write(
"\n==================================================================================================="
)
eval_results_logging_file.write("\n Test results for {} \n".format(
opt.model_name))
print(opt, file=eval_results_logging_file)
eval_results_logging_file.write(
"\n Test accuracy: Top1 {:.3f}, Top5 {:.3f}.".format(
test_acc, test_acc_top5))
eval_results_logging_file.write(
"\n Test error rate: Top1 {:.3f}, Top5 {:.3f}.".format(
100 - test_acc, 100 - test_acc_top5))