如何在PyTorch中无if分支程序化修改模型组件？

你完全不需要妥协于if分支或冗长的函数签名，有几种优雅的方案可以实现你的需求：

方案1：使用闭包绑定模型内部模块 #

通过定义返回闭包的工厂函数，让闭包捕获模型的内部模块（如linears），这样外部逻辑就能直接访问这些模块，同时保持forward方法简洁：

import torch
from typing import Callable

# 定义架构逻辑的工厂函数
def make_parallel_fn(linears):
    def fn(x1, x2):
        x1 = linears[0](x1)
        x2 = linears[1](x2)
        return x1 + x2
    return fn

def make_sequential_fn(linears):
    def fn(x1, x2):
        x = x1 + x2
        x = linears[0](x)
        x = linears[1](x)
        return x
    return fn

class Model(torch.nn.Module):
    def __init__(self, d_in: int, d_out: int, fn_builder: Callable):
        super().__init__()
        self.linears = torch.nn.ModuleList([
            torch.nn.Linear(d_in, d_out),
            torch.nn.Linear(d_in, d_out),
        ])
        # 传入内部模块生成绑定好的闭包
        self.forward_fn = fn_builder(self.linears)

    def forward(self, x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor:
        return self.forward_fn(x1, x2)

# 使用示例
parallel_model = Model(64, 32, make_parallel_fn)
sequential_model = Model(64, 32, make_sequential_fn)

方案2：将架构逻辑封装为子Module（推荐） #

这种方案最贴合PyTorch的设计理念，把不同的架构逻辑封装成独立的Module子类，主模型只需初始化对应的子模块即可，参数会被自动注册到主模型的优化器中：

import torch

# 定义不同架构的子模块
class ParallelBlock(torch.nn.Module):
    def __init__(self, linears):
        super().__init__()
        self.linears = linears
    
    def forward(self, x1, x2):
        x1 = self.linears[0](x1)
        x2 = self.linears[1](x2)
        return x1 + x2

class SequentialBlock(torch.nn.Module):
    def __init__(self, linears):
        super().__init__()
        self.linears = linears
    
    def forward(self, x1, x2):
        x = x1 + x2
        x = self.linears[0](x)
        x = self.linears[1](x)
        return x

class Model(torch.nn.Module):
    def __init__(self, d_in: int, d_out: int, block_cls):
        super().__init__()
        self.linears = torch.nn.ModuleList([
            torch.nn.Linear(d_in, d_out),
            torch.nn.Linear(d_in, d_out),
        ])
        # 初始化对应的架构子模块
        self.block = block_cls(self.linears)

    def forward(self, x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor:
        return self.block(x1, x2)

# 使用示例
parallel_model = Model(64, 32, ParallelBlock)
sequential_model = Model(64, 32, SequentialBlock)

方案3：使用functools.partial绑定参数 #

通过partial函数将模型内部模块预先绑定到外部函数的参数中，避免每次调用都手动传入：

import torch
from functools import partial
from typing import Callable

# 定义带模块参数的架构函数
def parallel_fn(linears, x1, x2):
    x1 = linears[0](x1)
    x2 = linears[1](x2)
    return x1 + x2

def sequential_fn(linears, x1, x2):
    x = x1 + x2
    x = linears[0](x)
    x = linears[1](x)
    return x

class Model(torch.nn.Module):
    def __init__(self, d_in: int, d_out: int, fn: Callable):
        super().__init__()
        self.linears = torch.nn.ModuleList([
            torch.nn.Linear(d_in, d_out),
            torch.nn.Linear(d_in, d_out),
        ])
        # 绑定linears为函数的第一个参数
        self.forward_fn = partial(fn, self.linears)

    def forward(self, x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor:
        return self.forward_fn(x1, x2)

# 使用示例
parallel_model = Model(64, 32, parallel_fn)
sequential_model = Model(64, 32, sequential_fn)

以上三种方案都能让你灵活切换模型架构逻辑，同时完全避免在forward方法中使用条件分支，也不需要把所有内部模块都传入外部函数。

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