ENAS代码解读

参考代码：https://github.com/TDeVries/enas_pytorch

数据集：cifar10

main函数：

def main():global argsnp.random.seed(args.seed)torch.cuda.manual_seed(args.seed)if args.fixed_arc:sys.stdout = Logger(filename='logs/' + args.output_filename + '_fixed.log')else:sys.stdout = Logger(filename='logs/' + args.output_filename + '.log')print(args)data_loaders = load_datasets()controller = Controller(search_for=args.search_for,search_whole_channels=True,num_layers=args.child_num_layers,num_branches=args.child_num_branches,out_filters=args.child_out_filters,lstm_size=args.controller_lstm_size,lstm_num_layers=args.controller_lstm_num_layers,tanh_constant=args.controller_tanh_constant,temperature=None,skip_target=args.controller_skip_target,skip_weight=args.controller_skip_weight)controller = controller.cuda()shared_cnn = SharedCNN(num_layers=args.child_num_layers,num_branches=args.child_num_branches,out_filters=args.child_out_filters,keep_prob=args.child_keep_prob)shared_cnn = shared_cnn.cuda()# https://github.com/melodyguan/enas/blob/master/src/utils.py#L218controller_optimizer = torch.optim.Adam(params=controller.parameters(),lr=args.controller_lr,betas=(0.0, 0.999),eps=1e-3)# https://github.com/melodyguan/enas/blob/master/src/utils.py#L213shared_cnn_optimizer = torch.optim.SGD(params=shared_cnn.parameters(),lr=args.child_lr_max,momentum=0.9,nesterov=True,weight_decay=args.child_l2_reg)# https://github.com/melodyguan/enas/blob/master/src/utils.py#L154shared_cnn_scheduler = CosineAnnealingLR(optimizer=shared_cnn_optimizer,T_max=args.child_lr_T,eta_min=args.child_lr_min)if args.resume:if os.path.isfile(args.resume):print("Loading checkpoint '{}'".format(args.resume))checkpoint = torch.load(args.resume)start_epoch = checkpoint['epoch']# args = checkpoint['args']shared_cnn.load_state_dict(checkpoint['shared_cnn_state_dict'])controller.load_state_dict(checkpoint['controller_state_dict'])shared_cnn_optimizer.load_state_dict(checkpoint['shared_cnn_optimizer'])controller_optimizer.load_state_dict(checkpoint['controller_optimizer'])shared_cnn_scheduler.optimizer = shared_cnn_optimizer  # Not sure if this actually worksprint("Loaded checkpoint '{}' (epoch {})".format(args.resume, checkpoint['epoch']))else:raise ValueError("No checkpoint found at '{}'".format(args.resume))else:start_epoch = 0if not args.fixed_arc:train_enas(start_epoch,controller,shared_cnn,data_loaders,shared_cnn_optimizer,controller_optimizer,shared_cnn_scheduler)else:assert args.resume != '', 'A pretrained model should be used when training a fixed architecture.'train_fixed(start_epoch,controller,shared_cnn,data_loaders)

再来看看Controller类的init

class Controller(nn.Module):'''https://github.com/melodyguan/enas/blob/master/src/cifar10/general_controller.py'''def __init__(self,search_for="macro",search_whole_channels=True,num_layers=12,num_branches=6,out_filters=36,lstm_size=32,lstm_num_layers=2,tanh_constant=1.5,temperature=None,skip_target=0.4,skip_weight=0.8):super(Controller, self).__init__()self.search_for = search_for # macroself.search_whole_channels = search_whole_channels # Trueself.num_layers = num_layers # 12self.num_branches = num_branches # 6self.out_filters = out_filters # 36self.lstm_size = lstm_size # 64self.lstm_num_layers = lstm_num_layers # 1self.tanh_constant = tanh_constant # 1.5self.temperature = temperature # Noneself.skip_target = skip_target # 0.4self.skip_weight = skip_weight # 0.8self._create_params()

num_layer为12代表最终生成12层的网络，num_branches为6代表6组操作：3x3,5x5正常卷积层，3x3,5x5深度分离卷积层，平均池化和最大池化，

再看Controller的 _create_params(self)函数：

def _create_params(self):'''https://github.com/melodyguan/enas/blob/master/src/cifar10/general_controller.py#L83'''self.w_lstm = nn.LSTM(input_size=self.lstm_size,hidden_size=self.lstm_size,num_layers=self.lstm_num_layers)self.g_emb = nn.Embedding(1, self.lstm_size)  # Learn the starting inputif self.search_whole_channels:self.w_emb = nn.Embedding(self.num_branches, self.lstm_size)self.w_soft = nn.Linear(self.lstm_size, self.num_branches, bias=False)else:assert False, "Not implemented error: search_whole_channels = False"self.w_attn_1 = nn.Linear(self.lstm_size, self.lstm_size, bias=False)self.w_attn_2 = nn.Linear(self.lstm_size, self.lstm_size, bias=False)self.v_attn = nn.Linear(self.lstm_size, 1, bias=False)self._reset_params()

没啥好解释的，这里值得注意的是第九行，controller的初始输入为1的embbeding:

self.g_emb = nn.Embedding(1, self.lstm_size)

重点看看forward()函数：

def forward(self):'''https://github.com/melodyguan/enas/blob/master/src/cifar10/general_controller.py#L126'''h0 = None  # setting h0 to None will initialize LSTM state with 0sanchors = []anchors_w_1 = []arc_seq = {}entropys = []log_probs = []skip_count = []skip_penaltys = []inputs = self.g_emb.weight# print('cccccccccccccccc')# print(inputs)# print(inputs.shape)# import sys# sys.exit()skip_targets = torch.tensor([1.0 - self.skip_target, self.skip_target]).cuda()for layer_id in range(self.num_layers):if self.search_whole_channels:inputs = inputs.unsqueeze(0)output, hn = self.w_lstm(inputs, h0)output = output.squeeze(0)h0 = hnlogit = self.w_soft(output)if self.temperature is not None:logit /= self.temperatureif self.tanh_constant is not None:logit = self.tanh_constant * torch.tanh(logit)branch_id_dist = Categorical(logits=logit)branch_id = branch_id_dist.sample()arc_seq[str(layer_id)] = [branch_id]log_prob = branch_id_dist.log_prob(branch_id)log_probs.append(log_prob.view(-1))entropy = branch_id_dist.entropy()entropys.append(entropy.view(-1))inputs = self.w_emb(branch_id)inputs = inputs.unsqueeze(0)else:# https://github.com/melodyguan/enas/blob/master/src/cifar10/general_controller.py#L171assert False, "Not implemented error: search_whole_channels = False"output, hn = self.w_lstm(inputs, h0)output = output.squeeze(0)if layer_id > 0:query = torch.cat(anchors_w_1, dim=0)query = torch.tanh(query + self.w_attn_2(output))query = self.v_attn(query)logit = torch.cat([-query, query], dim=1)if self.temperature is not None:logit /= self.temperatureif self.tanh_constant is not None:logit = self.tanh_constant * torch.tanh(logit)skip_dist = Categorical(logits=logit)skip = skip_dist.sample()skip = skip.view(layer_id)arc_seq[str(layer_id)].append(skip)skip_prob = torch.sigmoid(logit)kl = skip_prob * torch.log(skip_prob / skip_targets)kl = torch.sum(kl)skip_penaltys.append(kl)log_prob = skip_dist.log_prob(skip)log_prob = torch.sum(log_prob)log_probs.append(log_prob.view(-1))entropy = skip_dist.entropy()entropy = torch.sum(entropy)entropys.append(entropy.view(-1))# Calculate average hidden state of all nodes that got skips# and use it as input for next stepskip = skip.type(torch.float)skip = skip.view(1, layer_id)skip_count.append(torch.sum(skip))inputs = torch.matmul(skip, torch.cat(anchors, dim=0))inputs /= (1.0 + torch.sum(skip))else:inputs = self.g_emb.weightanchors.append(output)anchors_w_1.append(self.w_attn_1(output))self.sample_arc = arc_seqentropys = torch.cat(entropys)self.sample_entropy = torch.sum(entropys)log_probs = torch.cat(log_probs)self.sample_log_prob = torch.sum(log_probs)skip_count = torch.stack(skip_count)self.skip_count = torch.sum(skip_count)skip_penaltys = torch.stack(skip_penaltys)self.skip_penaltys = torch.mean(skip_penaltys)

32-35行及61至64行对应的是论文Training details里描述的：我们将tanh常数2.5和温度5.0应用于控制器的logits，并将控制器样本的熵添加到奖励中，权重为0.1。

37-38行是按logit的概率进行采样，返回采样的index。

42行log_prob解释（https://pytorch.org/docs/stable/distributions.html#torch.distributions.categorical.Categorical.log_prob）:

当概率密度函数的参数可微时，我们只需要sample()和log_prob()来实现REINFORCE:

在实践中，我们将从网络的输出中取样一个操作，将该操作应用到环境中，然后使用log_prob构造一个等效的损失函数。注意，我们使用了一个负数，因为优化器使用梯度下降，而上面的规则假设梯度上升。有了明确的政策，实施加强的准则如下:

probs = policy_network(state)
# Note that this is equivalent to what used to be called multinomial
m = Categorical(probs)
action = m.sample()
next_state, reward = env.step(action)
loss = -m.log_prob(action) * reward
loss.backward()

接下来42-45将控制器的样本熵添加到奖励中，然后把被采样的branch_id(对应于上面描述的6个卷积，池化等操作)，再把branch_id embbeding一下，得到下一时刻LSTM的输入。

72-75行对应的论文里Training details里描述的：在宏搜索空间中，我们通过增加两层之间的KL散度来增强跳跃连接的稀疏性：1）任意两层之间的跳跃连接概率；2）我们选择的概率ρ=0.4，它表示形成跳跃连接的先验信念。这个KL发散项的权重是0.8。

KL散度公式：

57-59行为attention。

60行是为了形成0，1分类，

接下来再看看ShareCNN( )代码：

先看 init(）函数：

class SharedCNN(nn.Module):def __init__(self,num_layers=12,num_branches=6,out_filters=24,keep_prob=1.0,fixed_arc=None):super(SharedCNN, self).__init__()self.num_layers = num_layers # 12self.num_branches = num_branches # 6self.out_filters = out_filters # 36self.keep_prob = keep_prob # 0.9self.fixed_arc = fixed_arcpool_distance = self.num_layers // 3self.pool_layers = [pool_distance - 1, 2 * pool_distance - 1]self.stem_conv = nn.Sequential(nn.Conv2d(3, out_filters, kernel_size=3, padding=1, bias=False),nn.BatchNorm2d(out_filters, track_running_stats=False))self.layers = nn.ModuleList([])self.pooled_layers = nn.ModuleList([])for layer_id in range(self.num_layers):if self.fixed_arc is None:layer = ENASLayer(layer_id, self.out_filters, self.out_filters)else:layer = FixedLayer(layer_id, self.out_filters, self.out_filters, self.fixed_arc[str(layer_id)])self.layers.append(layer)if layer_id in self.pool_layers:for i in range(len(self.layers)):if self.fixed_arc is None:self.pooled_layers.append(FactorizedReduction(self.out_filters, self.out_filters))else:self.pooled_layers.append(FactorizedReduction(self.out_filters, self.out_filters * 2))if self.fixed_arc is not None:self.out_filters *= 2self.global_avg_pool = nn.AdaptiveAvgPool2d((1, 1))self.dropout = nn.Dropout(p=1. - self.keep_prob)self.classify = nn.Linear(self.out_filters, 10)for m in self.modules():if isinstance(m, nn.Conv2d):nn.init.kaiming_uniform_(m.weight, mode='fan_in', nonlinearity='relu')

注意17-18行，第3，7层为pool layer。20-22行输入经过一个3x3卷积和BatchNorm层，再进入网络架构搜索模块。当fixed_arc为None时，调用ENASLayer()函数，代码如下：

class ENASLayer(nn.Module):'''https://github.com/melodyguan/enas/blob/master/src/cifar10/general_child.py#L245'''def __init__(self, layer_id, in_planes, out_planes):super(ENASLayer, self).__init__()self.layer_id = layer_idself.in_planes = in_planesself.out_planes = out_planesself.branch_0 = ConvBranch(in_planes, out_planes, kernel_size=3)self.branch_1 = ConvBranch(in_planes, out_planes, kernel_size=3, separable=True)self.branch_2 = ConvBranch(in_planes, out_planes, kernel_size=5)self.branch_3 = ConvBranch(in_planes, out_planes, kernel_size=5, separable=True)self.branch_4 = PoolBranch(in_planes, out_planes, 'avg')self.branch_5 = PoolBranch(in_planes, out_planes, 'max')self.bn = nn.BatchNorm2d(out_planes, track_running_stats=False)def forward(self, x, prev_layers, sample_arc):layer_type = sample_arc[0]if self.layer_id > 0:skip_indices = sample_arc[1]else:skip_indices = []if layer_type == 0:out = self.branch_0(x)elif layer_type == 1:out = self.branch_1(x)elif layer_type == 2:out = self.branch_2(x)elif layer_type == 3:out = self.branch_3(x)elif layer_type == 4:out = self.branch_4(x)elif layer_type == 5:out = self.branch_5(x)else:raise ValueError("Unknown layer_type {}".format(layer_type))for i, skip in enumerate(skip_indices):if skip == 1:out += prev_layers[i]out = self.bn(out)return out

branch_0 — branch_5分别为3x3正常卷积，3x3分离卷积，5x5正常卷积，5x5分离卷积，平均池化，最大池化。

回到ShareCNN的init函数，当layer_id为3，7时，后面还要加上一个Reduction层，及FactorizedReduction()函数。

FactorizedReduction函数代码：

class FactorizedReduction(nn.Module):'''Reduce both spatial dimensions (width and height) by a factor of 2, and potentially to change the number of output filtershttps://github.com/melodyguan/enas/blob/master/src/cifar10/general_child.py#L129'''def __init__(self, in_planes, out_planes, stride=2):super(FactorizedReduction, self).__init__()assert out_planes % 2 == 0, ("Need even number of filters when using this factorized reduction.")self.in_planes = in_planesself.out_planes = out_planesself.stride = strideif stride == 1:self.fr = nn.Sequential(nn.Conv2d(in_planes, out_planes, kernel_size=1, bias=False),nn.BatchNorm2d(out_planes, track_running_stats=False))else:self.path1 = nn.Sequential(nn.AvgPool2d(1, stride=stride),nn.Conv2d(in_planes, out_planes // 2, kernel_size=1, bias=False))self.path2 = nn.Sequential(nn.AvgPool2d(1, stride=stride),nn.Conv2d(in_planes, out_planes // 2, kernel_size=1, bias=False))self.bn = nn.BatchNorm2d(out_planes, track_running_stats=False)def forward(self, x):if self.stride == 1:return self.fr(x)else:path1 = self.path1(x)# pad the right and the bottom, then crop to include those pixelspath2 = F.pad(x, pad=(0, 1, 0, 1), mode='constant', value=0.)path2 = path2[:, :, 1:, 1:]path2 = self.path2(path2)out = torch.cat([path1, path2], dim=1)out = self.bn(out)return out

代码逻辑很简洁，没啥好介绍的。

接下来看train_eans( )函数：

def train_enas(start_epoch,controller,shared_cnn,data_loaders,shared_cnn_optimizer,controller_optimizer,shared_cnn_scheduler):"""Perform architecture search by training a controller and shared_cnn.Args:start_epoch: Epoch to begin on.controller: Controller module that generates architectures to be trained.shared_cnn: CNN that contains all possible architectures, with shared weights.data_loaders: Dict containing data loaders.shared_cnn_optimizer: Optimizer for the shared_cnn.controller_optimizer: Optimizer for the controller.shared_cnn_scheduler: Learning rate schedular for shared_cnn_optimizerReturns: Nothing."""baseline = Nonefor epoch in range(start_epoch, args.num_epochs):train_shared_cnn(epoch,controller,shared_cnn,data_loaders,shared_cnn_optimizer)baseline = train_controller(epoch,controller,shared_cnn,data_loaders,controller_optimizer,baseline)if epoch % args.eval_every_epochs == 0:evaluate_model(epoch, controller, shared_cnn, data_loaders)shared_cnn_scheduler.step(epoch)state = {'epoch': epoch + 1,'args': args,'shared_cnn_state_dict': shared_cnn.state_dict(),'controller_state_dict': controller.state_dict(),'shared_cnn_optimizer': shared_cnn_optimizer.state_dict(),'controller_optimizer': controller_optimizer.state_dict()}filename = 'checkpoints/' + args.output_filename + '.pth.tar'torch.save(state, filename)

先固定Controller，训练ShareCNN，再固定ShareCNN，训练Controller，迭代进行。

train_share_cnn()函数：

def train_shared_cnn(epoch,controller,shared_cnn,data_loaders,shared_cnn_optimizer,fixed_arc=None):"""Train shared_cnn by sampling architectures from the controller.Args:epoch: Current epoch.controller: Controller module that generates architectures to be trained.shared_cnn: CNN that contains all possible architectures, with shared weights.data_loaders: Dict containing data loaders.shared_cnn_optimizer: Optimizer for the shared_cnn.fixed_arc: Architecture to train, overrides the controller sample...Returns: Nothing."""global vis_wincontroller.eval()if fixed_arc is None:# Use a subset of the training set when searching for an arhcitecturetrain_loader = data_loaders['train_subset']else:# Use the full training set when training a fixed architecturetrain_loader = data_loaders['train_dataset']train_acc_meter = AverageMeter()loss_meter = AverageMeter()for i, (images, labels) in enumerate(train_loader):start = time.time()images = images.cuda()labels = labels.cuda()if fixed_arc is None:with torch.no_grad():controller()  # perform forward pass to generate a new architecturesample_arc = controller.sample_arcelse:sample_arc = fixed_arcshared_cnn.zero_grad()pred = shared_cnn(images, sample_arc)loss = nn.CrossEntropyLoss()(pred, labels)loss.backward()grad_norm = torch.nn.utils.clip_grad_norm_(shared_cnn.parameters(), args.child_grad_bound)shared_cnn_optimizer.step()train_acc = torch.mean((torch.max(pred, 1)[1] == labels).type(torch.float))train_acc_meter.update(train_acc.item())loss_meter.update(loss.item())end = time.time()if (i) % args.log_every == 0:learning_rate = shared_cnn_optimizer.param_groups[0]['lr']display = 'epoch=' + str(epoch) + \'\tch_step=' + str(i) + \'\tloss=%.6f' % (loss_meter.val) + \'\tlr=%.4f' % (learning_rate) + \'\t|g|=%.4f' % (grad_norm) + \'\tacc=%.4f' % (train_acc_meter.val) + \'\ttime=%.2fit/s' % (1. / (end - start))print(display)controller.train()

先用Controller sample出一个子模型，然后进行训练，然后记录acc。

train_controller代码：

def train_controller(epoch,controller,shared_cnn,data_loaders,controller_optimizer,baseline=None):"""Train controller to optimizer validation accuracy using REINFORCE.Args:epoch: Current epoch.controller: Controller module that generates architectures to be trained.shared_cnn: CNN that contains all possible architectures, with shared weights.data_loaders: Dict containing data loaders.controller_optimizer: Optimizer for the controller.baseline: The baseline score (i.e. average val_acc) from the previous epochReturns: baseline: The baseline score (i.e. average val_acc) for the current epochFor more stable training we perform weight updates using the average ofmany gradient estimates. controller_num_aggregate indicates how many sampleswe want to average over (default = 20). By default PyTorch will sum gradientseach time .backward() is called (as long as an optimizer step is not taken),so each iteration we divide the loss by controller_num_aggregate to get the average.https://github.com/melodyguan/enas/blob/master/src/cifar10/general_controller.py#L270"""print('Epoch ' + str(epoch) + ': Training controller')# global vis_winshared_cnn.eval()valid_loader = data_loaders['valid_subset']reward_meter = AverageMeter()baseline_meter = AverageMeter()val_acc_meter = AverageMeter()loss_meter = AverageMeter()controller.zero_grad()for i in range(args.controller_train_steps * args.controller_num_aggregate):start = time.time()images, labels = next(iter(valid_loader))images = images.cuda()labels = labels.cuda()controller()  # perform forward pass to generate a new architecturesample_arc = controller.sample_arcwith torch.no_grad():pred = shared_cnn(images, sample_arc)val_acc = torch.mean((torch.max(pred, 1)[1] == labels).type(torch.float))# detach to make sure that gradients aren't backpropped through the rewardreward = torch.tensor(val_acc.detach())reward += args.controller_entropy_weight * controller.sample_entropyif baseline is None:baseline = val_accelse:baseline -= (1 - args.controller_bl_dec) * (baseline - reward)# detach to make sure that gradients are not backpropped through the baselinebaseline = baseline.detach()loss = -1 * controller.sample_log_prob * (reward - baseline)if args.controller_skip_weight is not None:loss += args.controller_skip_weight * controller.skip_penaltysreward_meter.update(reward.item())baseline_meter.update(baseline.item())val_acc_meter.update(val_acc.item())loss_meter.update(loss.item())# Average gradient over controller_num_aggregate samplesloss = loss / args.controller_num_aggregateloss.backward(retain_graph=True)end = time.time()# Aggregate gradients for controller_num_aggregate iterationa, then update weightsif (i + 1) % args.controller_num_aggregate == 0:grad_norm = torch.nn.utils.clip_grad_norm_(controller.parameters(), args.child_grad_bound)controller_optimizer.step()controller.zero_grad()if (i + 1) % (2 * args.controller_num_aggregate) == 0:learning_rate = controller_optimizer.param_groups[0]['lr']display = 'ctrl_step=' + str(i // args.controller_num_aggregate) + \'\tloss=%.3f' % (loss_meter.val) + \'\tent=%.2f' % (controller.sample_entropy.item()) + \'\tlr=%.4f' % (learning_rate) + \'\t|g|=%.4f' % (grad_norm) + \'\tacc=%.4f' % (val_acc_meter.val) + \'\tbl=%.2f' % (baseline_meter.val) + \'\ttime=%.2fit/s' % (1. / (end - start))print(display)shared_cnn.train()return baseline

用强化学习训练，用采样子模型的acc和控制器熵添作为reward，熵权重为0.1，为了减少方差，reward减去一个baseline，baseline为reward的移动平均基线。然后用策略梯度公式更新controller。66行为计算强化学习loss。68-69行为把前面提到的KL作为损失加入进loss里，以增强跳跃连接的稀疏性。加个负号后反向传播loss，用梯度下降更新模型。

补充：

---

得分函数

当概率密度函数相对于其参数可微分时, 我们只需要sample()和log_prob()来实现REINFORCE:

是参数, 是学习速率, 是奖励 并且 是在状态 以及给定策略 执行动作 的概率.

在实践中, 我们将从网络输出中采样一个动作, 将这个动作应用于一个环境中, 然后使用log_prob构造一个等效的损失函数. 请注意, 我们使用负数是因为优化器使用梯度下降, 而上面的规则假设梯度上升. 有了确定的策略, REINFORCE的实现代码如下:

probs = policy_network(state)
# Note that this is equivalent to what used to be called multinomial
m = Categorical(probs)
action = m.sample()
next_state, reward = env.step(action)
loss = -m.log_prob(action) * reward
loss.backward()

转存中…(img-vAU4Xiu8-1594179710154)]

[外链图片转存中…(img-2SQqq9RE-1594179710155)] 是参数, [外链图片转存中…(img-XItub8OJ-1594179710156)] 是学习速率, [外链图片转存中…(img-7in5zxVt-1594179710156)] 是奖励 并且 [外链图片转存中…(img-0pFmPxGF-1594179710157)] 是在状态 [外链图片转存中…(img-hHi0NgtV-1594179710157)] 以及给定策略 [外链图片转存中…(img-lTprVpkz-1594179710157)]执行动作 [外链图片转存中…(img-x22cUOjD-1594179710158)] 的概率.

在实践中, 我们将从网络输出中采样一个动作, 将这个动作应用于一个环境中, 然后使用log_prob构造一个等效的损失函数. 请注意, 我们使用负数是因为优化器使用梯度下降, 而上面的规则假设梯度上升. 有了确定的策略, REINFORCE的实现代码如下:

probs = policy_network(state)
# Note that this is equivalent to what used to be called multinomial
m = Categorical(probs)
action = m.sample()
next_state, reward = env.step(action)
loss = -m.log_prob(action) * reward
loss.backward()

---