如何实现ResNet输出喂入EfficientNet的端到端反向传播优化？

我们有两个模型：自行实现的ResNet，以及通过torchvision.models库导入的EfficientNet。

ResNet的实现代码如下：

def __init__(self, in_channels, out_channels, num_classes):
    super().__init__()
    self.conv1_0 = _conv3x3(3, 32, stride=2)
    self.bn1_0 = _bn(32)
    self.conv1_1 = _conv3x3(32, 32, stride=1)
    self.bn1_1 = _bn(32)
    self.conv1_2 = _conv3x3(32, 64, stride=1)
    self.relu = nn.ReLU()
    self.pad = torch.nn.ReplicationPad2d(padding=(0, 0, 1, 1))
    self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2)
    self.end_point = _fc(225, num_classes) 

def forward(self, x):
    x = self.conv1(x)
    #and so on...     
    out = self.end_point(x)
    return out

EfficientNet的导入代码如下：

efficientNet = models.efficientnet_b5().to(device)

我们希望让EfficientNet接收ResNet的输出（中间用gumbel_model处理），并实现从EfficientNet输出到ResNet第一层的端到端优化。由于EfficientNet结构复杂，直接修改其forward方法比较困难。

目前我们仅基于EfficientNet的输出计算损失，优化器只传入EfficientNet的参数：

optimizer = optim.SGD([dict(params=efficientNet.parameters(), lr=LR)])

训练循环代码如下：

efficientNet = get_efficient_trained(device, out_features=1, in_features=3,path_model)
predictor = ResNet(ResidualBlockBase, layer_config, num_classes=num_cl)        

for i, imgs in enumerate(dataloader):
    inputs, labels = imgs
    inputs, labels = inputs.to(device), labels.to(device)  
    predictor_output = predictor(inputs)
    predictor_gumbel_output = gumbel(predictor_output)
    optimizer.zero_grad()
    outputs = efficientnet(predictor_gumbel_output).torch.squeeze(outputs, 1)     
    loss = loss_fn(outputs, labels)
    loss.backward()
    optimizer.step()    
return model

1. 当前是否仅对EfficientNet进行反向传播？
   是的。你的优化器只注册了EfficientNet的参数，训练循环中loss.backward()虽然会计算ResNet参数的梯度，但optimizer.step()只会更新EfficientNet的参数，ResNet的梯度不会被应用，相当于仅对EfficientNet做了反向传播优化。

2. 实现多模型端到端反向传播，是否需要将ResNet的参数加入优化器？若是，如何获取ResNet的参数？
   需要。端到端训练要求所有参与前向传播的可训练参数都被加入优化器，这样反向传播产生的梯度才能用来更新这些参数。
   获取ResNet参数直接调用predictor.parameters()即可，和获取EfficientNet参数的方式一致。

3. 如何让反向优化覆盖到ResNet，实现真正的端到端训练？
   修改优化器初始化代码，加入ResNet的参数，同时确保gumbel操作可微分，具体步骤：

• 初始化优化器时同时传入两个模型的参数（可单独设置学习率）：

optimizer = optim.SGD(
    [
        dict(params=efficientNet.parameters(), lr=LR),
        dict(params=predictor.parameters(), lr=LR)
    ]
)

• 保证gumbel函数可微分，比如使用Gumbel-Softmax时训练阶段设置hard=False，否则梯度无法传递到ResNet。
• 原训练循环无需大幅修改，loss.backward()会自动计算所有可训练参数的梯度，optimizer.step()会同时更新两个模型的参数，从而实现从EfficientNet输出到ResNet第一层的端到端优化。

内容的提问来源于stack exchange，提问作者Giovanna Purgato