使用torchsummary打印模型摘要时触发RuntimeError的问题求助

错误原因 #

nn.Embedding层要求输入必须是整数类型张量（LongTensor/IntTensor），它接收的是词汇表索引值而非浮点型数据。但torchsummary默认生成的输入张量是Float类型，导致类型不匹配触发错误。

解决方案 #

1. 修正torchsummary的输入张量类型 #

直接生成符合要求的整数类型输入张量，替代原来的形状参数传给torchsummary：

from torchsummary import summary
import torch

n_layers = 5
batch_size = 32
hidden_dim = 64
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model_v3 = SentimentalModelV3(output_size = 3, 
                              vocab_size = UNIQUE_WORD_COUNT, 
                              embedding_dim = 128, 
                              hidden_dim = hidden_dim, 
                              n_layers = n_layers, 
                              drop_prob = 0.3, 
                              padded_seq_len = 10, 
                              batch_size = batch_size,
                              bidirectional = False).to(device)

# 生成整数类型的输入索引张量，范围在0到词汇表大小之间
input_indices = torch.randint(0, UNIQUE_WORD_COUNT, (batch_size, 10)).to(device)
hn, cn = model_v3.initCellState()

# 传入具体的张量而非形状，确保类型匹配
summary(model_v3, [input_indices, hn, cn])

2. 优化模型代码（可选但更健壮） #

原模型存在几个潜在问题，建议同步修正：

• 硬编码batch_size会导致模型在推理时无法适配不同批次大小
• one_hot参数命名错误（Embedding接收的是索引，不是one-hot向量）
• 双向LSTM时，隐藏状态的层数需要乘以2

修改后的模型代码：

import torch.nn as nn
import torch

class SentimentalModelV3(nn.Module):
  def __init__(self, output_size, vocab_size, embedding_dim = 128, hidden_dim = 64, padded_seq_len = 10, n_layers = 1, drop_prob = 0.3, bidirectional = False):
    super().__init__()
    self.output_size = output_size
    self.n_layers = n_layers
    self.hidden_dim = hidden_dim
    self.padded_seq_len = padded_seq_len
    self.bidirectional = bidirectional

    self.embedding = nn.Embedding(num_embeddings = vocab_size, embedding_dim = embedding_dim)
    self.lstm = nn.LSTM(input_size = embedding_dim, hidden_size = hidden_dim, 
                        num_layers = n_layers, dropout = drop_prob, 
                        batch_first = True, bidirectional = bidirectional)

    self.dropout = nn.Dropout(0.3)

    # 适配双向LSTM的输出维度
    lstm_output_dim = hidden_dim * 2 if bidirectional else hidden_dim
    self.fc1=nn.Linear(lstm_output_dim * padded_seq_len, 64)
    self.fc2=nn.Linear(64, 16)
    self.fc3=nn.Linear(16,output_size)
    self.Relu = nn.ReLU()

  def forward(self, input_indices, hn, cn):
    embed = self.embedding(input_indices)
    lstm_out, hidden = self.lstm(embed, (hn, cn))

    # 展平LSTM输出
    lstm_out = lstm_out.reshape(lstm_out.shape[0], lstm_out.shape[1] * lstm_out.shape[2])
    out = self.dropout(lstm_out)
    out = self.Relu(out)
    out = self.Relu(self.fc1(out))
    out = self.Relu(self.fc2(out))
    out = self.fc3(out)

    return out

  def initCellState(self, batch_size):
    # 动态计算隐藏状态的层数和设备
    direction_mult = 2 if self.bidirectional else 1
    device = next(self.parameters()).device
    h = torch.zeros(self.n_layers * direction_mult, batch_size, self.hidden_dim).to(device)
    c = torch.zeros(self.n_layers * direction_mult, batch_size, self.hidden_dim).to(device)
    return h, c

对应调用代码调整为动态传入batch_size：

hn, cn = model_v3.initCellState(batch_size)
summary(model_v3, [input_indices, hn, cn])

关键说明 #

• torch.randint生成的张量默认是Long类型，完美匹配Embedding层的输入要求
• 移除硬编码的batch_size后，模型可以灵活适配不同的批次大小
• 双向LSTM的隐藏状态层数需要乘以2，否则会出现维度不匹配错误

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