def validate(run_manager, epoch=0, is_test=True, image_size_list=None,
width_mult_list=None, ks_list=None, expand_ratio_list=None, depth_list=None, additional_setting=None):
dynamic_net = run_manager.net
if isinstance(dynamic_net, nn.DataParallel):
dynamic_net = dynamic_net.module
dynamic_net.eval()
if image_size_list is None:
image_size_list = int2list(run_manager.run_config.data_provider.image_size, 1)
if width_mult_list is None:
width_mult_list = [i for i in range(len(dynamic_net.width_mult_list))]
if ks_list is None:
ks_list = dynamic_net.ks_list
if expand_ratio_list is None:
expand_ratio_list = dynamic_net.expand_ratio_list
if depth_list is None:
depth_list = dynamic_net.depth_list
subnet_settings = []
for w in width_mult_list:
for d in depth_list:
for e in expand_ratio_list:
for k in ks_list:
for img_size in image_size_list:
subnet_settings.append([{
'image_size': img_size,
'wid': w,
'd': d,
'e': e,
'ks': k,
}, 'R%s-W%s-D%s-E%s-K%s' % (img_size, w, d, e, k)])
if additional_setting is not None:
subnet_settings += additional_setting
losses_of_subnets, top1_of_subnets, top5_of_subnets = [], [], []
valid_log = ''
for setting, name in subnet_settings:
run_manager.write_log('-' * 30 + ' Validate %s ' % name + '-' * 30, 'train', should_print=False)
run_manager.run_config.data_provider.assign_active_img_size(setting.pop('image_size'))
dynamic_net.set_active_subnet(**setting)
run_manager.write_log(dynamic_net.module_str, 'train', should_print=False)
run_manager.reset_running_statistics(dynamic_net)
loss, top1, top5 = run_manager.validate(epoch=epoch, is_test=is_test, run_str=name, net=dynamic_net)
losses_of_subnets.append(loss)
top1_of_subnets.append(top1)
top5_of_subnets.append(top5)
valid_log += '%s (%.3f), ' % (name, top1)
return list_mean(losses_of_subnets), list_mean(top1_of_subnets), list_mean(top5_of_subnets), valid_log
def train(self, args, warmup_epochs=5, warmup_lr=0):
for epoch in range(self.start_epoch, self.run_config.n_epochs + warmup_epochs):
train_loss, train_top1, train_top5 = self.train_one_epoch(args, epoch, warmup_epochs, warmup_lr)
img_size, val_loss, val_top1, val_top5 = self.validate_all_resolution(epoch, is_test=False)
is_best = list_mean(val_top1) > self.best_acc
self.best_acc = max(self.best_acc, list_mean(val_top1))
if self.is_root:
val_log = '[{0}/{1}]\tloss {2:.3f}\ttop-1 acc {3:.3f} ({4:.3f})\ttop-5 acc {5:.3f}\t' \
'Train top-1 {top1:.3f}\tloss {train_loss:.3f}\t'. \
format(epoch + 1 - warmup_epochs, self.run_config.n_epochs, list_mean(val_loss),
list_mean(val_top1), self.best_acc, list_mean(val_top5),
top1=train_top1, train_loss=train_loss)
for i_s, v_a in zip(img_size, val_top1):
val_log += '(%d, %.3f), ' % (i_s, v_a)
self.write_log(val_log, prefix='valid', should_print=False)
self.save_model({
'epoch': epoch,
'best_acc': self.best_acc,
'optimizer': self.optimizer.state_dict(),
'state_dict': self.net.state_dict(),
}, is_best=is_best)
def train_one_epoch(run_manager, args, epoch, warmup_epochs=0, warmup_lr=0):
dynamic_net = run_manager.net
# switch to train mode
dynamic_net.train()
run_manager.run_config.train_loader.sampler.set_epoch(epoch)
MyRandomResizedCrop.EPOCH = epoch
nBatch = len(run_manager.run_config.train_loader)
data_time = AverageMeter()
losses = DistributedMetric('train_loss')
top1 = DistributedMetric('train_top1')
top5 = DistributedMetric('train_top5')
with tqdm(total=nBatch,
desc='Train Epoch #{}'.format(epoch + 1),
disable=not run_manager.is_root) as t:
end = time.time()
for i, (images,
labels) in enumerate(run_manager.run_config.train_loader):
data_time.update(time.time() - end)
if epoch < warmup_epochs:
new_lr = run_manager.run_config.warmup_adjust_learning_rate(
run_manager.optimizer,
warmup_epochs * nBatch,
nBatch,
epoch,
i,
warmup_lr,
)
else:
new_lr = run_manager.run_config.adjust_learning_rate(
run_manager.optimizer, epoch - warmup_epochs, i, nBatch)
images, labels = images.cuda(), labels.cuda()
target = labels
# soft target
if args.kd_ratio > 0:
args.teacher_model.train()
with torch.no_grad():
soft_logits = args.teacher_model(images).detach()
soft_label = F.softmax(soft_logits, dim=1)
# clear gradients
run_manager.optimizer.zero_grad()
loss_of_subnets, acc1_of_subnets, acc5_of_subnets = [], [], []
# compute output
subnet_str = ''
for _ in range(args.dynamic_batch_size):
# set random seed before sampling
if args.independent_distributed_sampling:
subnet_seed = os.getpid() + time.time()
else:
subnet_seed = int('%d%.3d%.3d' %
(epoch * nBatch + i, _, 0))
random.seed(subnet_seed)
subnet_settings = dynamic_net.sample_active_subnet()
subnet_str += '%d: ' % _ + ','.join([
'%s_%s' %
(key, '%.1f' %
subset_mean(val, 0) if isinstance(val, list) else val)
for key, val in subnet_settings.items()
]) + ' || '
output = run_manager.net(images)
if args.kd_ratio == 0:
loss = run_manager.train_criterion(output, labels)
loss_type = 'ce'
else:
if args.kd_type == 'ce':
kd_loss = cross_entropy_loss_with_soft_target(
output, soft_label)
else:
kd_loss = F.mse_loss(output, soft_logits)
loss = args.kd_ratio * kd_loss + run_manager.train_criterion(
output, labels)
loss = loss * (2 / (args.kd_ratio + 1))
loss_type = '%.1fkd-%s & ce' % (args.kd_ratio,
args.kd_type)
# measure accuracy and record loss
acc1, acc5 = accuracy(output, target, topk=(1, 5))
loss_of_subnets.append(loss)
acc1_of_subnets.append(acc1[0])
acc5_of_subnets.append(acc5[0])
loss.backward()
run_manager.optimizer.step()
losses.update(list_mean(loss_of_subnets), images.size(0))
top1.update(list_mean(acc1_of_subnets), images.size(0))
top5.update(list_mean(acc5_of_subnets), images.size(0))
t.set_postfix({
'loss': losses.avg.item(),
'top1': top1.avg.item(),
'top5': top5.avg.item(),
'R': images.size(2),
'lr': new_lr,
'loss_type': loss_type,
'seed': str(subnet_seed),
'str': subnet_str,
'data_time': data_time.avg,
})
t.update(1)
end = time.time()
return losses.avg.item(), top1.avg.item(), top5.avg.item()
def validate(run_manager,
epoch=0,
is_test=True,
image_size_list=None,
width_mult_list=None,
ks_list=None,
expand_ratio_list=None,
depth_list=None,
pixelshuffle_depth_list=None,
additional_setting=None):
dynamic_net = run_manager.net
if isinstance(dynamic_net, nn.DataParallel):
dynamic_net = dynamic_net.module
dynamic_net.eval()
if image_size_list is None:
image_size_list = int2list(
run_manager.run_config.data_provider.image_size, 1)
if width_mult_list is None:
width_mult_list = [i for i in range(len(dynamic_net.width_mult_list))]
if ks_list is None:
ks_list = dynamic_net.ks_list
if expand_ratio_list is None:
expand_ratio_list = dynamic_net.expand_ratio_list
if depth_list is None:
depth_list = dynamic_net.depth_list
if pixelshuffle_depth_list is None:
pixelshuffle_depth_list = dynamic_net.pixelshuffle_depth_list
subnet_settings = []
for pixel_d in pixelshuffle_depth_list:
for w in width_mult_list:
for d in depth_list:
for e in expand_ratio_list:
for k in ks_list:
# for img_size in image_size_list:
subnet_settings.append([
{
# 'image_size': img_size,
'pixel_d': pixel_d,
'wid': w,
'd': d,
'e': e,
'ks': k,
},
'PD%s-W%s-D%s-E%s-K%s' % (pixel_d, w, d, e, k)
])
if additional_setting is not None:
subnet_settings += additional_setting
# losses_of_subnets, top1_of_subnets, top5_of_subnets = [], [], []
losses_of_subnets, psnr_of_subnets = [], []
valid_log = ''
for setting, name in subnet_settings:
#################### Validation Architecture 정하는 부분인데, Single Architecture Overfitting 혹은 뭐 빠르게 테스트 해볼일 있으면 여기서 그냥 스킵하면됨
# if name.find('PD1-W0-D2-E3-K7') == -1:
# continue
run_manager.write_log('-' * 30 + ' Validate %s ' % name + '-' * 30,
'train',
should_print=False)
# run_manager.run_config.data_provider.assign_active_img_size(setting.pop('image_size'))
#################### Random Sampling과 Structured Sampling중에 주석 바꿔가면서 고르면 됨
# dynamic_net.sample_active_subnet()
dynamic_net.set_active_subnet(**setting)
run_manager.write_log(dynamic_net.module_str,
'train',
should_print=False)
#################### Oracle Training 시에는 Batch Mean/Variance 현재 데이터로 업데이트하면 망함.
# run_manager.reset_running_statistics(dynamic_net)
loss, psnr = run_manager.validate(epoch=epoch,
is_test=is_test,
run_str=name,
net=dynamic_net)
losses_of_subnets.append(loss)
# top1_of_subnets.append(top1)
# top5_of_subnets.append(top5)
psnr_of_subnets.append(psnr)
valid_log += '%s (%.3f), ' % (name, psnr)
return list_mean(losses_of_subnets), list_mean(psnr_of_subnets), valid_log
def train_one_epoch(run_manager, args, epoch, warmup_epochs=0, warmup_lr=0):
dynamic_net = run_manager.net
if isinstance(dynamic_net, nn.DataParallel):
dynamic_net = dynamic_net.module
# switch to train mode
dynamic_net.train()
# run_manager.run_config.train_loader.sampler.set_epoch(epoch)
# MyRandomResizedCrop.EPOCH = epoch
#################### Code for freezing BN. Overfitting 할 때는 주석 해제하면됨.
# for m in dynamic_net.modules():
# if isinstance(m, nn.BatchNorm2d):
# ########## Use running mean/var
# m.eval()
# ########## BN weight/bias freeze
# # m.weight.requires_grad = False
# # m.bias.requires_grad = False
nBatch = len(run_manager.run_config.train_loader)
data_time = AverageMeter()
# losses = DistributedMetric('train_loss')
# top1 = DistributedMetric('train_top1')
# top5 = DistributedMetric('train_top5')
losses = AverageMeter()
psnr_averagemeter = AverageMeter()
with tqdm(total=nBatch, desc='Train Epoch #{}'.format(epoch + 1)) as t:
end = time.time()
for i, mini_batch in enumerate(run_manager.run_config.train_loader):
images = mini_batch['image']
#################### 2x or 4x 고르는 부분.
x2_down_images = mini_batch['2x_down_image']
x4_down_images = mini_batch['4x_down_image']
data_time.update(time.time() - end)
if epoch < warmup_epochs:
new_lr = run_manager.run_config.warmup_adjust_learning_rate(
run_manager.optimizer,
warmup_epochs * nBatch,
nBatch,
epoch,
i,
warmup_lr,
)
else:
new_lr = run_manager.run_config.adjust_learning_rate(
run_manager.optimizer, epoch - warmup_epochs, i, nBatch)
images = images.cuda()
#################### 2x or 4x 고르는 부분.
x2_down_images = x2_down_images.cuda()
x4_down_images = x4_down_images.cuda()
target = images
# soft target
if args.kd_ratio > 0:
args.teacher_model.train()
with torch.no_grad():
soft_logits = args.teacher_model(images).detach()
soft_label = F.softmax(soft_logits, dim=1)
# clear gradients
run_manager.optimizer.zero_grad()
loss_of_subnets, psnr_of_subnets = [], []
# compute output
subnet_str = ''
for _ in range(args.dynamic_batch_size):
# set random seed before sampling
if args.independent_distributed_sampling:
subnet_seed = os.getpid() + time.time()
else:
subnet_seed = int('%d%.3d%.3d' %
(epoch * nBatch + i, _, 0))
random.seed(subnet_seed)
#################### Random Sampling과 Structured Sampling중에 주석 바꿔가면서 고르면 됨. Single Architecture Overfitting을 위해서 여기 수정해주면 가능.
subnet_settings = dynamic_net.sample_active_subnet()
# dynamic_net.set_active_subnet(ks=7, e=3, d=2, pixel_d=1)
subnet_str += '%d: ' % _ + ','.join([
'%s_%s' %
(key, '%.1f' %
subset_mean(val, 0) if isinstance(val, list) else val)
for key, val in subnet_settings.items()
]) + ' || '
#################### 2x or 4x 고르는 부분.
# output = run_manager.net(images)
if subnet_settings['pixel_d'][0] == 1:
output = run_manager.net(x2_down_images)
elif subnet_settings['pixel_d'][0] == 2:
output = run_manager.net(x4_down_images)
if args.kd_ratio == 0:
loss = run_manager.train_criterion(output, images)
loss_type = 'mse'
else:
if args.kd_type == 'ce':
kd_loss = cross_entropy_loss_with_soft_target(
output, soft_label)
else:
kd_loss = F.mse_loss(output, soft_logits)
loss = args.kd_ratio * kd_loss + run_manager.train_criterion(
output, labels)
loss = loss * (2 / (args.kd_ratio + 1))
loss_type = '%.1fkd-%s & ce' % (args.kd_ratio,
args.kd_type)
# measure accuracy and record loss
# acc1, acc5 = accuracy(output, target, topk=(1, 5))
psnr_current = psnr(rgb2y(tensor2img_np(output)),
rgb2y(tensor2img_np(images)))
loss_of_subnets.append(loss)
# acc1_of_subnets.append(acc1[0])
# acc5_of_subnets.append(acc5[0])
psnr_of_subnets.append(psnr_current)
loss.backward()
run_manager.optimizer.step()
losses.update(list_mean(loss_of_subnets), images.size(0))
# top1.update(list_mean(acc1_of_subnets), images.size(0))
# top5.update(list_mean(acc5_of_subnets), images.size(0))
psnr_averagemeter.update(list_mean(psnr_of_subnets),
images.size(0))
t.set_postfix({
'loss': losses.avg.item(),
# 'top1': top1.avg.item(),
# 'top5': top5.avg.item(),
'psnr': psnr_averagemeter.avg,
'R': images.size(2),
'lr': new_lr,
'loss_type': loss_type,
'seed': str(subnet_seed),
'str': subnet_str,
'data_time': data_time.avg,
})
t.update(1)
end = time.time()
return losses.avg.item(), psnr_averagemeter.avg