PyTorch C++中GRUCell的SeLU激活函数替换实现问题

问题背景 #

任务目标是修改PyTorch源码中GRUCell组件的new_gate计算逻辑（对应原代码第799行），将tanh_()替换为SeLU激活函数。原代码片段如下：

template <typename cell_params>
struct GRUCell : Cell<Tensor, cell_params> {
  using hidden_type = Tensor;

  hidden_type operator()(
      const Tensor& input,
      const hidden_type& hidden,
      const cell_params& params,
      bool pre_compute_input = false) const override {
    if (input.is_cuda() || input.is_xpu()) {
      TORCH_CHECK(!pre_compute_input);
      auto igates = params.matmul_ih(input);
      auto hgates = params.matmul_hh(hidden);
      auto result = at::_thnn_fused_gru_cell(
          igates, hgates, hidden, params.b_ih(), params.b_hh());
      // Slice off the workspace argument (it's needed only for AD).
      return std::move(std::get<0>(result));
    }
    const auto chunked_igates = pre_compute_input
        ? input.unsafe_chunk(3, 1)
        : params.linear_ih(input).unsafe_chunk(3, 1);
    auto chunked_hgates = params.linear_hh(hidden).unsafe_chunk(3, 1);
    const auto reset_gate =
        chunked_hgates[0].add_(chunked_igates[0]).sigmoid_();
    const auto input_gate =
        chunked_hgates[1].add_(chunked_igates[1]).sigmoid_();
    const auto new_gate =
        chunked_igates[2].add(chunked_hgates[2].mul_(reset_gate)).tanh_();
    return (hidden - new_gate).mul_(input_gate).add_(new_gate);
  }
};

尝试直接替换tanh_()为selu_()并引入头文件后，编译报错“Did you mean relu_()”；自定义selu函数时又遇到Tensor数据类型不兼容问题。

解决方案 #

1. 优先使用PyTorch内置SeLU函数（规避in-place限制） #

PyTorch C++ API中，selu_()（in-place版本）并非对所有Tensor类型/设备后端都支持，优先用非in-place的selu()函数替代：

// 替换原new_gate行
const auto new_gate = chunked_igates[2].add(chunked_hgates[2].mul_(reset_gate)).selu();

需确保引入头文件：

#include <ATen/ATen.h>

2. 自定义兼容多类型的SeLU实现（解决类型兼容问题） #

如果内置函数仍有类型适配问题，手动实现SeLU的数学逻辑，利用Tensor的apply()方法遍历元素：
SeLU公式：
$$ \text{SeLU}(x) = \lambda \times \begin{cases} x & x > 0 \ \alpha (e^x - 1) & x \leq 0 \end{cases} $$
默认参数$\lambda=1.0507$，$\alpha=1.67326$。

替换代码如下：

// 先计算new_gate的输入值
const auto new_gate_input = chunked_igates[2].add(chunked_hgates[2].mul_(reset_gate));
// 自定义SeLU处理，自动适配Tensor数据类型
auto new_gate = new_gate_input.apply([](const auto& x) {
    const auto alpha = static_cast<decltype(x)>(1.67326);
    const auto lambda = static_cast<decltype(x)>(1.0507);
    if (x > 0) {
        return lambda * x;
    } else {
        return lambda * alpha * (std::exp(x) - static_cast<decltype(x)>(1.0));
    }
});

apply()方法会自动处理Tensor的每个元素，适配float、double等不同数据类型，同时支持CPU/CUDA设备。

3. 处理CUDA/XPU分支逻辑 #

原代码中CUDA/XPU路径使用了at::_thnn_fused_gru_cell融合操作，该函数内部硬编码了tanh激活。若需全平台替换，需修改该融合函数的底层实现；若仅需CPU环境生效，保持CUDA分支逻辑不变即可。

内容的提问来源于stack exchange，提问作者Tahseen Abbas Hyder