如何将PyTorch nn.Module与Cython结合优化DeepSet网络前向传播？

优先尝试：PyTorch批量处理替代Python循环 #

你的核心瓶颈是Python循环中逐个调用pers_lay1，可以通过批量合并所有子集合元素的方式，利用PyTorch的C++级批量运算替代循环，这比Cython实现更简单且效率更高：

class DSSN(nn.Module):
    def __init__(self, pd_rho, dh_rho, fc_network, device):
        super(DSSN, self).__init__()
        self.pers_lay1 = pd_rho
        self.pers_lay2 = dh_rho
        self.fc = fc_network

    def forward(self, x: List[List[torch.Tensor]]) -> torch.Tensor:
        # 1. 扁平化所有子集合元素，组成批量
        all_pd = [pd for sample in x for pd in sample]
        batch_pd = torch.stack(all_pd)
        
        # 2. 批量通过pers_lay1，替代逐个调用
        processed_batch = self.pers_lay1(batch_pd)
        
        # 3. 将批量结果拆分回原嵌套结构
        processed_samples = []
        idx = 0
        for sample in x:
            num_pd = len(sample)
            processed_samples.append(processed_batch[idx:idx+num_pd])
            idx += num_pd
        
        # 4. 后续步骤保持不变
        x = torch.stack([torch.stack(sample) for sample in processed_samples])
        x = self.pers_lay2(x)
        x = self.fc(x)
        return x

这种方式将多次pers_lay1调用合并为一次批量运算，完全利用PyTorch的优化计算图，避免Python循环的开销，速度提升通常远超过TorchScript的40%。

若批量处理不可行（元素形状/长度不一致）：Cython结合PyTorch方案 #

如果子集合元素形状或长度无法统一，再考虑用Cython优化循环。核心思路是将嵌套循环逻辑移到Cython中，减少Python解释器的开销：

步骤1：编写Cython工具函数 #

创建ds_utils.pyx文件：

import torch
cimport torch
from torch cimport Tensor

def process_nested_input(list[list[Tensor]] x, torch.nn.Module pers_lay1):
    cdef list processed_samples = []
    cdef list sample_processed
    cdef list sample
    cdef Tensor pd, pd_processed

    for sample in x:
        sample_processed = []
        for pd in sample:
            pd_processed = pers_lay1(pd)
            sample_processed.append(pd_processed)
        processed_samples.append(torch.stack(sample_processed))
    return torch.stack(processed_samples)

步骤2：编译Cython模块 #

创建setup.py：

from setuptools import setup, Extension
from Cython.Build import cythonize
import torch

extensions = [
    Extension(
        "ds_utils",
        ["ds_utils.pyx"],
        include_dirs=[torch.utils.cpp_extension.include_paths()[0]],
        language="c++",
    )
]

setup(
    name="ds_utils",
    ext_modules=cythonize(extensions),
)

执行编译命令：

python setup.py build_ext --inplace

步骤3：在PyTorch模块中调用Cython函数 #

from ds_utils import process_nested_input

class DSSN(nn.Module):
    def __init__(self, pd_rho, dh_rho, fc_network, device):
        super(DSSN, self).__init__()
        self.pers_lay1 = pd_rho
        self.pers_lay2 = dh_rho
        self.fc = fc_network

    def forward(self, x: List[List[torch.Tensor]]) -> torch.Tensor:
        x = process_nested_input(x, self.pers_lay1)
        x = self.pers_lay2(x)
        x = self.fc(x)
        return x

进阶优化：绕过Python层调用nn.Module #

如果pers_lay1是简单层（如线性层），可以直接在Cython中调用PyTorch的底层张量操作，彻底避免Python交互开销：

import torch
cimport torch
from torch cimport Tensor

def process_nested_input_linear(list[list[Tensor]] x, Tensor weight, Tensor bias):
    cdef list processed_samples = []
    cdef list sample_processed
    cdef list sample
    cdef Tensor pd, pd_processed

    for sample in x:
        sample_processed = []
        for pd in sample:
            # 直接执行线性变换，无需调用Python层的nn.Module
            pd_processed = torch.matmul(pd, weight.t()) + bias
            sample_processed.append(pd_processed)
        processed_samples.append(torch.stack(sample_processed))
    return torch.stack(processed_samples)

在forward中调用时，直接传入层的权重和偏置：

def forward(self, x: List[List[torch.Tensor]]) -> torch.Tensor:
    x = process_nested_input_linear(x, self.pers_lay1.weight, self.pers_lay1.bias)
    x = self.pers_lay2(x)
    x = self.fc(x)
    return x

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