训练pytorch模型后如何丢弃分支

Question

我正在尝试在pytorch 中实现一个 FCN，其整体结构如下：

到目前为止的代码如下所示：

class SNet(nn.Module):
    def __init__(self):
        super(SNet, self).__init__()
        
        self.enc_a = encoder(...)
        self.dec_a = decoder(...)
        
        self.enc_b = encoder(...)
        self.dec_b = decoder(...)
    
    def forward(self, x1, x2):
        x1 = self.enc_a(x1)
        x2 = self.enc_b(x2)
        x2 = self.dec_b(x2)
        x1 = self.dec_a(torch.cat((x1, x2), dim=-1)
        return x1, x2

在keras 中，使用函数式API 相对容易做到这一点。但是，我在pytorch 中找不到任何具体的示例/教程。

1. 如何在训练后丢弃dec_a（自动编码器分支的解码器部分）？
2. 在联合训练期间，loss 将是来自两个分支的loss 的总和（可选加权）？

Answer 1

您还可以为模型定义单独的模式以进行训练和推理：

class SNet(nn.Module):
  def __init__(self):
    super(SNet, self).__init__()
    
    self.enc_a = encoder(...)
    self.dec_a = decoder(...)
    
    self.enc_b = encoder(...)
    self.dec_b = decoder(...)
    
    self.training = True

  def forward(self, x1, x2):
    if self.training:
        x1 = self.enc_a(x1)
        x2 = self.enc_b(x2)
        x2 = self.dec_b(x2)
        x1 = self.dec_a(torch.cat((x1, x2), dim=-1)
        return x1, x2
    else:
        x1 = self.enc_a(x1)
        x2 = self.enc_b(x2)
        x2 = self.dec_b(x2)
        return x2

这些块是示例，可能无法完全按照您的意愿执行，因为我认为您在块图中定义训练和推理操作的方式与您的代码之间存在一些歧义，但无论如何您都明白了仅在训练模式下如何使用某些模块。然后你可以相应地设置这个变量。