python实现GRU并使用torch.nn.GRU验证正确性

最新推荐文章于 2024-08-13 21:32:11 发布

阿袁的小园子

最新推荐文章于 2024-08-13 21:32:11 发布

阅读量1.2k

点赞数

分类专栏： python 算法文章标签： python

本文链接：https://blog.csdn.net/yuanren201/article/details/124422346

版权

python 同时被 2 个专栏收录

36 篇文章 3 订阅

订阅专栏

算法

14 篇文章 0 订阅

订阅专栏

这篇博客展示了如何使用Python基础语法和PyTorch库手动实现门控循环单元（GRU），并将其结果与PyTorch内置的GRU模块进行对比。通过allclose函数验证，两者计算结果一致，证明手动实现的GRU是正确的。

摘要由CSDN通过智能技术生成

用python基本语法以及torch的一些基本数据结构对GRU进行实现如下：

import torch
import torch.nn as nn

bs,T,i_size,h_size=2,3,4,5

input=torch.rand(bs,T,i_size)

h_0=torch.rand(bs,h_size)

gru=nn.GRU(i_size,h_size,batch_first=True)
output,h_n=gru(input,h_0.unsqueeze(0))
print(output)
print(h_n)

def gru_forward(input,initial_states,w_ih,w_hh,b_ih,b_hh):
    bs,T,i_size=input.shape
    h_size=initial_states.shape[-1]
    h_0=initial_states

    batch_w_ih=w_ih.unsqueeze(0).tile(bs,1,1)
    batch_w_hh=w_hh.unsqueeze(0).tile(bs,1,1)


    prev_h=h_0

    for t in range(T):
        x=input[:,t,:]
        w_times_x=torch.bmm(batch_w_ih,x.unsqueeze(-1)).squeeze(-1)
        w_times_h=torch.bmm(batch_w_hh,prev_h.unsqueeze(-1)).squeeze(-1)
        r=torch.sigmoid(w_times_x[:,:h_size]+w_times_h[:,:h_size]+b_ih[:h_size]+b_hh[:h_size])
        z=torch.sigmoid(w_times_x[:,h_size:2*h_size]+w_times_h[:,h_size:2*h_size]+\
            b_ih[h_size:2*h_size]+b_hh[h_size:2*h_size])
        n=torch.tanh(w_times_x[:,2*h_size:3*h_size]+b_ih[2*h_size:3*h_size]+\
            r*(w_times_h[:,2*h_size:3*h_size]+b_hh[2*h_size:3*h_size]))
        prev_h=(1-z)*n+z*prev_h
        output[:,t,:]=prev_h
    return output,prev_h

output_custom,h_n_custom=gru_forward(input,h_0,gru.weight_ih_l0,gru.weight_hh_l0,gru.bias_ih_l0,gru.bias_hh_l0)

print(torch.allclose(output,output_custom))
print(torch.allclose(h_n,h_n_custom))