pytorch 调用lstm

module

import torch
import torch.nn as nn

class M(nn.Module):
	def __init__(self):
		super().__init__()
		self.lstm = nn.LSTM(3, 3, 1)  # input's dim = 3, hidden'dim = 3, num of lstm = 1
	def forward(self, x):
		out = self.lstm(*x)
		return out

data

data = torch.randn(2, 1, 3)  # seq_len=5, batch_size=1, dim=3; each epoch get 1 sentence, with per sentence have 2 words.
h_data = torch.randn(1, 1, 3)  # 1 lstm layer, 1 batch_size, 3 hidden node
c_data = torch.randn(1, 1, 3)

input_data = (data, h_data, c_data)

test

module = M()
output, (h_out, c_out) = module(input_data)

猜你喜欢

转载自blog.csdn.net/qq_38973721/article/details/120900227