conv_lstm实现

conv_lstm实现
看大佬们的代码，各种调用，读起来有点不方便，所以写一个特别简单的
首先要写两个类，一个conv_lstm_cell，一个conv_lstm
要实现的功能，把一个3通道的3030的图像卷价成1通道3030的图像，lstm层数是5（其实是几在这个小实验里头关系不大）

import torch
import torch.nn as nn
from torch.autograd import Variable
import numpy as np

class ConvLSTM_Cell(nn.Module):

def __init__(self):
    super(ConvLSTM_Cell, self).__init__()
    self.conv = nn.Conv2d(in_channels=4, out_channels=1, padding=1, kernel_size=3)
    
def forward(self, x, h_prev):
    combine = torch.cat((x, h_prev), dim=1)
    conv = self.conv(combine)
    return conv

class ConvLSTM(nn.Module):

def __init__(self):
    super(ConvLSTM, self).__init__()
    cell_list=[]
    for i in range(5):
        cell_list.append(ConvLSTM_Cell())
    self.cell_list=nn.ModuleList(cell_list)

def forward(self, x, hidden_state=None):
    if hidden_state is None:
        h = self.init_hidden()

    for layeridx in range(5):
        h = self.cell_list[layeridx](x, h_prev=h)
    return h
def init_hidden(self):
    return torch.FloatTensor(np.random.rand(1,1,30,30))

if name == ‘main’:
net=ConvLSTM()
mnist_img=Variable(torch.FloatTensor(np.random.rand(1,3,30,30)))
h=net(mnist_img)
print(h.size())