def main():
global args
args = parser.parse_args()
if args.cuda:
cudnn.benchmark = True
python_name = load_setting_aug.load_python_name(args, Path(__file__).name)
print(python_name)
global save_max_accu
save_max_accu = './{}.txt'.format(python_name)
# with open(save_max_accu, 'a') as f:
# f.write(python_name)
# f.write('\n')
# f.close()
for TIMES in range(1, args.TIMES):
nets, optimizers = load_setting_aug.define_net_and_opt(args)
# exit()
train_loader1, train_loader2, test_loader = load_setting_aug.load_dataset(
args)
max_prec = [0.0, 0.0, 0.0, 0.0, 0.0]
for epoch in range(1, args.epochs + 1):
epoch_start_time = time.time()
adjust_lr(optimizers, epoch, times=TIMES)
if using_one_train:
train(train_loader1, nets, optimizers, epoch)
else:
train1(train_loader1, nets, optimizers, epoch)
print(
'--------------------------------------------------------------------------'
)
train2(train_loader2, nets, optimizers, epoch)
epoch_time = time.time() - epoch_start_time
print('one epoch time is {:02}h{:02}m{:02}s'.format(
*transform_time(epoch_time)))
print('testing the models......')
pred_list = test12(test_loader, nets, epoch)
for id in range(len(pred_list)):
if max_prec[id] < pred_list[id]:
max_prec[id] = pred_list[id]
current_max = 'Net:{} || Current-Max:[{:.2f}, {:.2f}, {:.2f}, {:.2f}, {:.2f}]'.format(
python_name, max_prec[0], max_prec[1], max_prec[2],
max_prec[3], max_prec[4])
print(current_max)
if epoch == args.epochs:
with open(save_max_accu, 'a') as f:
f.write(str(current_max))
f.write('\n')
f.close()
def main():
global args
args = parser.parse_args()
print(args)
if not os.path.exists(os.path.join(args.save_root,'checkpoint')):
os.makedirs(os.path.join(args.save_root,'checkpoint'))
if args.cuda:
cudnn.benchmark = True
print('----------- Network Initialization --------------')
snet = define_tsnet(name=args.s_name, num_class=args.num_class, cuda=args.cuda)
checkpoint = torch.load(args.s_init)
load_pretrained_model(snet, checkpoint['net'])
tnet = define_tsnet(name=args.t_name, num_class=args.num_class, cuda=args.cuda)
checkpoint = torch.load(args.t_model)
load_pretrained_model(tnet, checkpoint['net'])
tnet.eval()
for param in tnet.parameters():
param.requires_grad = False
print('-----------------------------------------------')
# initialize optimizer
optimizer = torch.optim.SGD(snet.parameters(),
lr = args.lr,
momentum = args.momentum,
weight_decay = args.weight_decay,
nesterov = True)
# define loss functions
if args.cuda:
criterionCls = torch.nn.CrossEntropyLoss().cuda()
criterionFitnet = torch.nn.MSELoss().cuda()
else:
criterionCls = torch.nn.CrossEntropyLoss()
criterionFitnet = torch.nn.MSELoss()
# define transforms
if args.data_name == 'cifar10':
dataset = dst.CIFAR10
mean = (0.4914, 0.4822, 0.4465)
std = (0.2470, 0.2435, 0.2616)
elif args.data_name == 'cifar100':
dataset = dst.CIFAR100
mean = (0.5071, 0.4865, 0.4409)
std = (0.2673, 0.2564, 0.2762)
else:
raise Exception('invalid dataset name...')
train_transform = transforms.Compose([
transforms.Pad(4, padding_mode='reflect'),
transforms.RandomCrop(32),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize(mean=mean,std=std)
])
test_transform = transforms.Compose([
transforms.CenterCrop(32),
transforms.ToTensor(),
transforms.Normalize(mean=mean,std=std)
])
# define data loader
train_loader = torch.utils.data.DataLoader(
dataset(root = args.img_root,
transform = train_transform,
train = True,
download = True),
batch_size=args.batch_size, shuffle=True, num_workers=4, pin_memory=True)
test_loader = torch.utils.data.DataLoader(
dataset(root = args.img_root,
transform = test_transform,
train = False,
download = True),
batch_size=args.batch_size, shuffle=False, num_workers=4, pin_memory=True)
for epoch in range(1, args.epochs+1):
epoch_start_time = time.time()
adjust_lr(optimizer, epoch)
# train one epoch
nets = {'snet':snet, 'tnet':tnet}
criterions = {'criterionCls':criterionCls, 'criterionFitnet':criterionFitnet}
train(train_loader, nets, optimizer, criterions, epoch)
epoch_time = time.time() - epoch_start_time
print('one epoch time is {:02}h{:02}m{:02}s'.format(*transform_time(epoch_time)))
# evaluate on testing set
print('testing the models......')
test_start_time = time.time()
test(test_loader, nets, criterions)
test_time = time.time() - test_start_time
print('testing time is {:02}h{:02}m{:02}s'.format(*transform_time(test_time)))
# save model
print('saving models......')
save_name = 'fitnet_r{}_r{}_{:>03}.ckp'.format(args.t_name[6:], args.s_name[6:], epoch)
save_name = os.path.join(args.save_root, 'checkpoint', save_name)
if epoch == 1:
save_checkpoint({
'epoch': epoch,
'snet': snet.state_dict(),
'tnet': tnet.state_dict(),
}, save_name)
else:
save_checkpoint({
'epoch': epoch,
'snet': snet.state_dict(),
}, save_name)
def main():
global args
args = parser.parse_args()
print(args)
if not os.path.exists(os.path.join(args.save_root,'checkpoint')):
os.makedirs(os.path.join(args.save_root,'checkpoint'))
if args.cuda:
cudnn.benchmark = True
print('----------- Network Initialization --------------')
net = define_tsnet(name=args.net_name, num_class=args.num_class, cuda=args.cuda)
print('-----------------------------------------------')
# save initial parameters
print('saving initial parameters......')
if args.net_name[:6] == 'resnet':
save_name = 'baseline_r{}_{:>03}.ckp'.format(args.net_name[6:], 0)
elif args.net_name[:6] == 'resnex':
save_name = 'baseline_rx{}_{:>03}.ckp'.format(args.net_name[7:], 0)
elif args.net_name[:6] == 'densen':
save_name = 'baseline_dBC{}_{:>03}.ckp'.format(args.net_name[10:], 0)
save_name = os.path.join(args.save_root, 'checkpoint', save_name)
save_checkpoint({
'epoch': 0,
'net': net.state_dict(),
}, save_name)
# initialize optimizer
optimizer = torch.optim.SGD(net.parameters(),
lr = args.lr,
momentum = args.momentum,
weight_decay = args.weight_decay,
nesterov = True)
# define loss functions
if args.cuda:
criterion = torch.nn.CrossEntropyLoss().cuda()
else:
criterion = torch.nn.CrossEntropyLoss()
# define transforms
if args.data_name == 'cifar10':
dataset = dst.CIFAR10
mean = (0.4914, 0.4822, 0.4465)
std = (0.2470, 0.2435, 0.2616)
elif args.data_name == 'cifar100':
dataset = dst.CIFAR100
mean = (0.5071, 0.4865, 0.4409)
std = (0.2673, 0.2564, 0.2762)
else:
raise Exception('invalid dataset name...')
train_transform = transforms.Compose([
transforms.Pad(4, padding_mode='reflect'),
transforms.RandomCrop(32),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize(mean=mean,std=std)
])
test_transform = transforms.Compose([
transforms.CenterCrop(32),
transforms.ToTensor(),
transforms.Normalize(mean=mean,std=std)
])
# define data loader
train_loader = torch.utils.data.DataLoader(
dataset(root = args.img_root,
transform = train_transform,
train = True,
download = True),
batch_size=args.batch_size, shuffle=True, num_workers=4, pin_memory=True)
test_loader = torch.utils.data.DataLoader(
dataset(root = args.img_root,
transform = test_transform,
train = False,
download = True),
batch_size=args.batch_size, shuffle=False, num_workers=4, pin_memory=True)
trainF = open(os.path.join(args.save_log, 'train.csv'), 'w')
testF = open(os.path.join(args.save_log, 'test.csv'), 'w')
train_start_time = time.time()
max_test_prec_1, max_test_prec_5 = 0, 0
for epoch in range(1, args.epochs+1):
epoch_start_time = time.time()
adjust_lr(optimizer, epoch)
# train one epoch
train(train_loader, net, optimizer, criterion, epoch, trainF)
epoch_time = time.time() - epoch_start_time
print('one epoch time is {:02}h{:02}m{:02}s'.format(*transform_time(epoch_time)))
# evaluate on testing set
print('testing the models......')
test_start_time = time.time()
test_prec_1, test_prec_5 = test(test_loader, net, criterion, testF)
test_time = time.time() - test_start_time
print('testing time is {:02}h{:02}m{:02}s'.format(*transform_time(test_time)))
max_test_prec_1 = max(max_test_prec_1, test_prec_1)
max_test_prec_5 = max(max_test_prec_5, test_prec_5)
# save model
print('saving models......')
save_name = 'baseline_r{}_{:>03}.ckp'.format(args.net_name[6:], epoch)
save_name = os.path.join(args.save_root, 'checkpoint', save_name)
save_checkpoint({
'epoch': epoch,
'net': net.state_dict(),
}, save_name)
train_finish_time = time.time()
train_total_time = train_start_time - train_finish_time
trainF.close()
testF.close()
print('the total train time is:{:02}h{:02}m{:02}s'.format(*transform_time(train_total_time)))
print('the max test prec@1:{:.2f} , prec@5:{:.2f}'.format(max_test_prec_1, max_test_prec_5))
