From df45e68366f37cb9b4a1835da78fa34d8958d487 Mon Sep 17 00:00:00 2001 From: D-X-Y <280835372@qq.com> Date: Sat, 25 Jul 2020 12:39:55 +0000 Subject: [PATCH] Update TuNAS --- exps/algos-v2/search-size.py | 125 +++++++----------- exps/experimental/vis-bench-ws.py | 6 +- .../generic_size_tiny_cell_model.py | 19 ++- lib/nas_201_api/api_301.py | 1 + 4 files changed, 64 insertions(+), 87 deletions(-) diff --git a/exps/algos-v2/search-size.py b/exps/algos-v2/search-size.py index 1eebd4e..b43971f 100644 --- a/exps/algos-v2/search-size.py +++ b/exps/algos-v2/search-size.py @@ -8,6 +8,10 @@ # python ./exps/algos-v2/search-size.py --dataset cifar10 --data_path $TORCH_HOME/cifar.python --algo fbv2 --rand_seed 777 # python ./exps/algos-v2/search-size.py --dataset cifar100 --data_path $TORCH_HOME/cifar.python --algo fbv2 --rand_seed 777 # python ./exps/algos-v2/search-size.py --dataset ImageNet16-120 --data_path $TORCH_HOME/cifar.python/ImageNet16 --algo fbv2 --rand_seed 777 +#### +# python ./exps/algos-v2/search-size.py --dataset cifar10 --data_path $TORCH_HOME/cifar.python --algo tunas --arch_weight_decay 0 --rand_seed 777 --use_api 0 +# python ./exps/algos-v2/search-size.py --dataset cifar100 --data_path $TORCH_HOME/cifar.python --algo tunas --arch_weight_decay 0 --rand_seed 777 +# python ./exps/algos-v2/search-size.py --dataset ImageNet16-120 --data_path $TORCH_HOME/cifar.python/ImageNet16 --algo tunas --arch_weight_decay 0 --rand_seed 777 ###################################################################################### import os, sys, time, random, argparse import numpy as np @@ -26,7 +30,28 @@ from models import get_cell_based_tiny_net, get_search_spaces from nas_201_api import NASBench301API as API -def search_func(xloader, network, criterion, scheduler, w_optimizer, a_optimizer, epoch_str, print_freq, logger): +# Ad-hoc for TuNAS +class ExponentialMovingAverage(object): + """Class that maintains an exponential moving average.""" + + def __init__(self, momentum): + self._numerator = 0 + self._denominator = 0 + self._momentum = momentum + + def update(self, value): + self._numerator = self._momentum * self._numerator + (1 - self._momentum) * value + self._denominator = self._momentum * self._denominator + (1 - self._momentum) + + @property + def value(self): + """Return the current value of the moving average""" + return self._numerator / self._denominator + +RL_BASELINE_EMA = ExponentialMovingAverage(0.95) + + +def search_func(xloader, network, criterion, scheduler, w_optimizer, a_optimizer, algo, epoch_str, print_freq, logger): data_time, batch_time = AverageMeter(), AverageMeter() base_losses, base_top1, base_top5 = AverageMeter(), AverageMeter(), AverageMeter() arch_losses, arch_top1, arch_top5 = AverageMeter(), AverageMeter(), AverageMeter() @@ -43,7 +68,7 @@ def search_func(xloader, network, criterion, scheduler, w_optimizer, a_optimizer # Update the weights network.zero_grad() - _, logits = network(base_inputs) + _, logits, _ = network(base_inputs) base_loss = criterion(logits, base_targets) base_loss.backward() w_optimizer.step() @@ -55,12 +80,21 @@ def search_func(xloader, network, criterion, scheduler, w_optimizer, a_optimizer # update the architecture-weight network.zero_grad() - _, logits = network(arch_inputs) - arch_loss = criterion(logits, arch_targets) + _, logits, log_probs = network(arch_inputs) + arch_prec1, arch_prec5 = obtain_accuracy(logits.data, arch_targets.data, topk=(1, 5)) + if algo == 'tunas': + with torch.no_grad(): + RL_BASELINE_EMA.update(arch_prec1.item()) + rl_advantage = arch_prec1 - RL_BASELINE_EMA.value + rl_log_prob = sum(log_probs) + arch_loss = - rl_advantage * rl_log_prob + elif algo == 'tas' or algo == 'fbv2': + arch_loss = criterion(logits, arch_targets) + else: + raise ValueError('invalid algorightm name: {:}'.format(algo)) arch_loss.backward() a_optimizer.step() # record - arch_prec1, arch_prec5 = obtain_accuracy(logits.data, arch_targets.data, topk=(1, 5)) arch_losses.update(arch_loss.item(), arch_inputs.size(0)) arch_top1.update (arch_prec1.item(), arch_inputs.size(0)) arch_top5.update (arch_prec5.item(), arch_inputs.size(0)) @@ -78,76 +112,6 @@ def search_func(xloader, network, criterion, scheduler, w_optimizer, a_optimizer return base_losses.avg, base_top1.avg, base_top5.avg, arch_losses.avg, arch_top1.avg, arch_top5.avg -def train_controller(xloader, network, criterion, optimizer, prev_baseline, epoch_str, print_freq, logger): - # config. (containing some necessary arg) - # baseline: The baseline score (i.e. average val_acc) from the previous epoch - data_time, batch_time = AverageMeter(), AverageMeter() - GradnormMeter, LossMeter, ValAccMeter, EntropyMeter, BaselineMeter, RewardMeter, xend = AverageMeter(), AverageMeter(), AverageMeter(), AverageMeter(), AverageMeter(), AverageMeter(), time.time() - - controller_num_aggregate = 20 - controller_train_steps = 50 - controller_bl_dec = 0.99 - controller_entropy_weight = 0.0001 - - network.eval() - network.controller.train() - network.controller.zero_grad() - loader_iter = iter(xloader) - for step in range(controller_train_steps * controller_num_aggregate): - try: - inputs, targets = next(loader_iter) - except: - loader_iter = iter(xloader) - inputs, targets = next(loader_iter) - inputs = inputs.cuda(non_blocking=True) - targets = targets.cuda(non_blocking=True) - # measure data loading time - data_time.update(time.time() - xend) - - log_prob, entropy, sampled_arch = network.controller() - with torch.no_grad(): - network.set_cal_mode('dynamic', sampled_arch) - _, logits = network(inputs) - val_top1, val_top5 = obtain_accuracy(logits.data, targets.data, topk=(1, 5)) - val_top1 = val_top1.view(-1) / 100 - reward = val_top1 + controller_entropy_weight * entropy - if prev_baseline is None: - baseline = val_top1 - else: - baseline = prev_baseline - (1 - controller_bl_dec) * (prev_baseline - reward) - - loss = -1 * log_prob * (reward - baseline) - - # account - RewardMeter.update(reward.item()) - BaselineMeter.update(baseline.item()) - ValAccMeter.update(val_top1.item()*100) - LossMeter.update(loss.item()) - EntropyMeter.update(entropy.item()) - - # Average gradient over controller_num_aggregate samples - loss = loss / controller_num_aggregate - loss.backward(retain_graph=True) - - # measure elapsed time - batch_time.update(time.time() - xend) - xend = time.time() - if (step+1) % controller_num_aggregate == 0: - grad_norm = torch.nn.utils.clip_grad_norm_(network.controller.parameters(), 5.0) - GradnormMeter.update(grad_norm) - optimizer.step() - network.controller.zero_grad() - - if step % print_freq == 0: - Sstr = '*Train-Controller* ' + time_string() + ' [{:}][{:03d}/{:03d}]'.format(epoch_str, step, controller_train_steps * controller_num_aggregate) - Tstr = 'Time {batch_time.val:.2f} ({batch_time.avg:.2f}) Data {data_time.val:.2f} ({data_time.avg:.2f})'.format(batch_time=batch_time, data_time=data_time) - Wstr = '[Loss {loss.val:.3f} ({loss.avg:.3f}) Prec@1 {top1.val:.2f} ({top1.avg:.2f}) Reward {reward.val:.2f} ({reward.avg:.2f})] Baseline {basel.val:.2f} ({basel.avg:.2f})'.format(loss=LossMeter, top1=ValAccMeter, reward=RewardMeter, basel=BaselineMeter) - Estr = 'Entropy={:.4f} ({:.4f})'.format(EntropyMeter.val, EntropyMeter.avg) - logger.log(Sstr + ' ' + Tstr + ' ' + Wstr + ' ' + Estr) - - return LossMeter.avg, ValAccMeter.avg, BaselineMeter.avg, RewardMeter.avg - - def valid_func(xloader, network, criterion, logger): data_time, batch_time = AverageMeter(), AverageMeter() arch_losses, arch_top1, arch_top5 = AverageMeter(), AverageMeter(), AverageMeter() @@ -159,7 +123,7 @@ def valid_func(xloader, network, criterion, logger): # measure data loading time data_time.update(time.time() - end) # prediction - _, logits = network(arch_inputs.cuda(non_blocking=True)) + _, logits, _ = network(arch_inputs.cuda(non_blocking=True)) arch_loss = criterion(logits, arch_targets) # record arch_prec1, arch_prec5 = obtain_accuracy(logits.data, arch_targets.data, topk=(1, 5)) @@ -211,9 +175,9 @@ def main(xargs): params = count_parameters_in_MB(search_model) logger.log('The parameters of the search model = {:.2f} MB'.format(params)) logger.log('search-space : {:}'.format(search_space)) - try: + if bool(xargs.use_api): api = API(verbose=False) - except: + else: api = None logger.log('{:} create API = {:} done'.format(time_string(), api)) @@ -250,7 +214,7 @@ def main(xargs): network.set_tau( xargs.tau_max - (xargs.tau_max-xargs.tau_min) * epoch / (total_epoch-1) ) logger.log('[RESET tau as : {:}]'.format(network.tau)) search_w_loss, search_w_top1, search_w_top5, search_a_loss, search_a_top1, search_a_top5 \ - = search_func(search_loader, network, criterion, w_scheduler, w_optimizer, a_optimizer, epoch_str, xargs.print_freq, logger) + = search_func(search_loader, network, criterion, w_scheduler, w_optimizer, a_optimizer, xargs.algo, epoch_str, xargs.print_freq, logger) search_time.update(time.time() - start_time) logger.log('[{:}] search [base] : loss={:.2f}, accuracy@1={:.2f}%, accuracy@5={:.2f}%, time-cost={:.1f} s'.format(epoch_str, search_w_loss, search_w_top1, search_w_top5, search_time.sum)) logger.log('[{:}] search [arch] : loss={:.2f}, accuracy@1={:.2f}%, accuracy@5={:.2f}%'.format(epoch_str, search_a_loss, search_a_top1, search_a_top5)) @@ -305,8 +269,9 @@ if __name__ == '__main__': parser.add_argument('--data_path' , type=str, help='Path to dataset') parser.add_argument('--dataset' , type=str, choices=['cifar10', 'cifar100', 'ImageNet16-120'], help='Choose between Cifar10/100 and ImageNet-16.') parser.add_argument('--search_space', type=str, default='sss', choices=['sss'], help='The search space name.') - parser.add_argument('--algo' , type=str, choices=['tas', 'fbv2', 'enas'], help='The search space name.') + parser.add_argument('--algo' , type=str, choices=['tas', 'fbv2', 'tunas'], help='The search space name.') parser.add_argument('--genotype' , type=str, default='|nor_conv_3x3~0|+|nor_conv_3x3~0|nor_conv_3x3~1|+|skip_connect~0|nor_conv_3x3~1|nor_conv_3x3~2|', help='The genotype.') + parser.add_argument('--use_api' , type=int, default=1, choices=[0,1], help='Whether use API or not (which will cost much memory).') # FOR GDAS parser.add_argument('--tau_min', type=float, default=0.1, help='The minimum tau for Gumbel Softmax.') parser.add_argument('--tau_max', type=float, default=10, help='The maximum tau for Gumbel Softmax.') diff --git a/exps/experimental/vis-bench-ws.py b/exps/experimental/vis-bench-ws.py index ebb7739..7f055b0 100644 --- a/exps/experimental/vis-bench-ws.py +++ b/exps/experimental/vis-bench-ws.py @@ -29,8 +29,8 @@ from log_utils import time_string def fetch_data(root_dir='./output/search', search_space='tss', dataset=None): ss_dir = '{:}-{:}'.format(root_dir, search_space) alg2name, alg2path = OrderedDict(), OrderedDict() - seeds = [777] if search_space == 'tss': + seeds = [777] alg2name['GDAS'] = 'gdas-affine0_BN0-None' alg2name['RSPS'] = 'random-affine0_BN0-None' alg2name['DARTS (1st)'] = 'darts-v1-affine0_BN0-None' @@ -38,8 +38,10 @@ def fetch_data(root_dir='./output/search', search_space='tss', dataset=None): alg2name['ENAS'] = 'enas-affine0_BN0-None' alg2name['SETN'] = 'setn-affine0_BN0-None' else: + seeds = [777, 888, 999] alg2name['TAS'] = 'tas-affine0_BN0' alg2name['FBNetV2'] = 'fbv2-affine0_BN0' + alg2name['TuNAS'] = 'tunas-affine0_BN0' for alg, name in alg2name.items(): alg2path[alg] = os.path.join(ss_dir, dataset, name, 'seed-{:}-last-info.pth') alg2data = OrderedDict() @@ -84,7 +86,7 @@ def visualize_curve(api, vis_save_dir, search_space): alg2data = fetch_data(search_space=search_space, dataset=dataset) alg2accuracies = OrderedDict() epochs = 100 - colors = ['b', 'g', 'c', 'm', 'y'] + colors = ['b', 'g', 'c', 'm', 'y', 'r'] ax.set_xlim(0, epochs) # ax.set_ylim(y_min_s[(dataset, search_space)], y_max_s[(dataset, search_space)]) for idx, (alg, data) in enumerate(alg2data.items()): diff --git a/lib/models/shape_searchs/generic_size_tiny_cell_model.py b/lib/models/shape_searchs/generic_size_tiny_cell_model.py index 0996597..1cd4328 100644 --- a/lib/models/shape_searchs/generic_size_tiny_cell_model.py +++ b/lib/models/shape_searchs/generic_size_tiny_cell_model.py @@ -47,10 +47,10 @@ class GenericNAS301Model(nn.Module): def set_algo(self, algo: Text): # used for searching assert self._algo is None, 'This functioin can only be called once.' - assert algo in ['fbv2', 'enas', 'tas'], 'invalid algo : {:}'.format(algo) + assert algo in ['fbv2', 'tunas', 'tas'], 'invalid algo : {:}'.format(algo) self._algo = algo self._arch_parameters = nn.Parameter(1e-3*torch.randn(self._max_num_Cs, len(self._candidate_Cs))) - if algo == 'fbv2' or algo == 'enas': + if algo == 'fbv2' or algo == 'tunas': self.register_buffer('_masks', torch.zeros(len(self._candidate_Cs), max(self._candidate_Cs))) for i in range(len(self._candidate_Cs)): self._masks.data[i, :self._candidate_Cs[i]] = 1 @@ -106,15 +106,17 @@ class GenericNAS301Model(nn.Module): def forward(self, inputs): feature = inputs + + log_probs = [] for i, cell in enumerate(self._cells): feature = cell(feature) + # apply different searching algorithms + idx = max(0, i-1) if self._algo == 'fbv2': - idx = max(0, i-1) weights = nn.functional.gumbel_softmax(self._arch_parameters[idx:idx+1], tau=self.tau, dim=-1) mask = torch.matmul(weights, self._masks).view(1, -1, 1, 1) feature = feature * mask elif self._algo == 'tas': - idx = max(0, i-1) selected_cs, selected_probs = select2withP(self._arch_parameters[idx:idx+1], self.tau, num=2) with torch.no_grad(): i1, i2 = selected_cs.cpu().view(-1).tolist() @@ -128,6 +130,13 @@ class GenericNAS301Model(nn.Module): else: miss = torch.zeros(feature.shape[0], feature.shape[1]-out.shape[1], feature.shape[2], feature.shape[3], device=feature.device) feature = torch.cat((out, miss), dim=1) + elif self._algo == 'tunas': + prob = nn.functional.softmax(self._arch_parameters[idx:idx+1], dim=-1) + dist = torch.distributions.Categorical(prob) + action = dist.sample() + log_probs.append(dist.log_prob(action)) + mask = self._masks[action.item()].view(1, -1, 1, 1) + feature = feature * mask else: raise ValueError('invalid algorithm : {:}'.format(self._algo)) @@ -136,4 +145,4 @@ class GenericNAS301Model(nn.Module): out = out.view(out.size(0), -1) logits = self.classifier(out) - return out, logits + return out, logits, log_probs diff --git a/lib/nas_201_api/api_301.py b/lib/nas_201_api/api_301.py index 8ac77f8..005dc05 100644 --- a/lib/nas_201_api/api_301.py +++ b/lib/nas_201_api/api_301.py @@ -60,6 +60,7 @@ class NASBench301API(NASBenchMetaAPI): self.reset_time() if file_path_or_dict is None: file_path_or_dict = os.path.join(os.environ['TORCH_HOME'], ALL_BENCHMARK_FILES[-1]) + print ('Try to use the default NAS-Bench-301 path from {:}.'.format(file_path_or_dict)) if isinstance(file_path_or_dict, str) or isinstance(file_path_or_dict, Path): file_path_or_dict = str(file_path_or_dict) if verbose: print('try to create the NAS-Bench-201 api from {:}'.format(file_path_or_dict))