PyTorch用torchsummary打印含GRU的Sequential模型报ReLU输入错误如何解决

• PyTorch中nn.GRU层的前向传播返回值是二元组(output, h_n)，并非单个张量：
  • output是所有时间步的输出特征，形状为(batch_size, seq_len, hidden_size)
  • h_n是最后一个时间步的隐藏状态，形状为(num_layers, batch_size, hidden_size)
• 直接将nn.GRU放在nn.Sequential的第一层时，它返回的元组会直接传递给下一层nn.ReLU，而ReLU仅支持张量输入，因此触发类型错误。

提供两种可直接运行的修复方案：

方案1：新增适配层兼容Sequential写法 #

如果要保留nn.Sequential的写法，新增一个极简的自定义层提取GRU的有效输出即可：

import torch
import torch.nn as nn
import torch.nn.functional as F
from torchsummary import summary

# 自定义适配层：提取GRU的输出张量，丢弃不需要的隐藏状态
class GRUOutputExtractor(nn.Module):
    def forward(self, x):
        # 若业务需要使用最后一步隐藏状态，返回x[1]并对应调整后续层维度即可
        return x[0]

feature_extractor_1 = nn.Sequential(
        nn.GRU(8, 64, num_layers=2, batch_first=True),
        GRUOutputExtractor(), # 仅需新增这一行适配GRU输出格式
        nn.ReLU(True),
        nn.Linear(64, 64),
        nn.ReLU(True)
    ).to('cuda')

summary(feature_extractor_1, input_size=(2500,8), device='cuda')

方案2：自定义nn.Module类（更推荐） #

对于包含循环层的模型，自定义模型类的写法灵活度更高，后续修改逻辑也更方便：

import torch
import torch.nn as nn
import torch.nn.functional as F
from torchsummary import summary

class FeatureExtractor(nn.Module):
    def __init__(self):
        super().__init__()
        self.gru = nn.GRU(8, 64, num_layers=2, batch_first=True)
        self.relu1 = nn.ReLU(True)
        self.linear1 = nn.Linear(64, 64)
        self.relu2 = nn.ReLU(True)
    
    def forward(self, x):
        # 显式解包GRU返回值，丢弃不需要的隐藏状态
        x, _ = self.gru(x)
        x = self.relu1(x)
        x = self.linear1(x)
        x = self.relu2(x)
        return x

feature_extractor_1 = FeatureExtractor().to('cuda')
summary(feature_extractor_1, input_size=(2500,8), device='cuda')

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