用pytorch写GRU

可以使用以下代码来使用 PyTorch 编写 GRU 模型：

import torch
import torch.nn as nn

class GRUModel(nn.Module):
def __init__(self, input_size, hidden_size, num_layers, output_size):
super(GRUModel, self).__init__()

self.hidden_size = hidden_size
self.num_layers = num_layers

self.gru = nn.GRU(input_size, hidden_size, num_layers, batch_first=True)
self.fc = nn.Linear(hidden_size, output_size)

def forward(self, x, h):
out, h = self.gru(x, h)
out = self.fc(out[:, -1, :])
return out, h

def init_hidden(self, batch_size):
return torch.zeros(self.num_layers, batch_size, self.hidden_size)

# 创建模型
input_size = 10
hidden_size = 20
num_layers = 2
output_size = 1

model = GRUModel(input_size, hidden_size, num_layers, output_size)

# 定义损失函数和优化器
criterion = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=.001)

# 训练模型
for epoch in range(num_epochs):
for i, (inputs, labels) in enumerate(train_loader):
# 初始化隐藏状态
h = model.init_hidden(batch_size)

# 前向传播
outputs, h = model(inputs, h)
loss = criterion(outputs, labels)

# 反向传播和优化
optimizer.zero_grad()
loss.backward()
optimizer.step()

# 打印日志
if (i+1) % 100 == :
print('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}'
.format(epoch+1, num_epochs, i+1, total_step, loss.item()))