怎么直接调用pytorch中的convlstm

你可以使用以下代码来直接调用PyTorch中的ConvLSTM：

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

class ConvLSTMCell(nn.Module):
    def __init__(self, input_size, hidden_size, kernel_size):
        super(ConvLSTMCell, self).__init__()
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.kernel_size = kernel_size
        self.padding = kernel_size // 2

        self.Wxi = nn.Conv2d(in_channels=input_size, out_channels=hidden_size, kernel_size=kernel_size, padding=self.padding)
        self.Whi = nn.Conv2d(in_channels=hidden_size, out_channels=hidden_size, kernel_size=kernel_size, padding=self.padding)
        self.Wxf = nn.Conv2d(in_channels=input_size, out_channels=hidden_size, kernel_size=kernel_size, padding=self.padding)
        self.Whf = nn.Conv2d(in_channels=hidden_size, out_channels=hidden_size, kernel_size=kernel_size, padding=self.padding)
        self.Wxc = nn.Conv2d(in_channels=input_size, out_channels=hidden_size, kernel_size=kernel_size, padding=self.padding)
        self.Whc = nn.Conv2d(in_channels=hidden_size, out_channels=hidden_size, kernel_size=kernel_size, padding=self.padding)
        self.Wxo = nn.Conv2d(in_channels=input_size, out_channels=hidden_size, kernel_size=kernel_size, padding=self.padding)
        self.Who = nn.Conv2d(in_channels=hidden_size, out_channels=hidden_size, kernel_size=kernel_size, padding=self.padding)

        self.Wci = None
        self.Wcf = None
        self.Wco = None

    def forward(self, x, h, c):
        ci = torch.sigmoid(self.Wxi(x) + self.Whi(h) + c * self.Wci)
        cf = torch.sigmoid(self.Wxf(x) + self.Whf(h) + c * self.Wcf)
        cc = cf * c + ci * torch.tanh(self.Wxc(x) + self.Whc(h))
        co = torch.sigmoid(self.Wxo(x) + self.Who(h) + cc * self.Wco)
        ch = co * torch.tanh(cc)

        return ch, cc

class ConvLSTM(nn.Module):
    def __init__(self, input_size, hidden_size, kernel_size, num_layers, batch_first=False):
        super(ConvLSTM, self).__init__()

        self.input_size = input_size
        self.hidden_size = hidden_size
        self.kernel_size = kernel_size
        self.num_layers = num_layers
        self.batch_first = batch_first

        cell_list = []
        for i in range(0, self.num_layers):
            cur_input_size = self.input_size if i == 0 else self.hidden_size
            cell_list.append(ConvLSTMCell(input_size=cur_input_size, hidden_size=self.hidden_size, kernel_size=self.kernel_size))

        self.cell_list = nn.ModuleList(cell_list)

    def forward(self, input_tensor, hidden_state=None):
        if hidden_state is None:
            hidden_state = self._init_hidden(batch_size=input_tensor.size(0))

        layer_output_list = []
        last_state_list = []

        seq_len = input_tensor.size(1) if self.batch_first else input_tensor.size(0)

        cur_layer_input = input_tensor

        for layer_idx in range(self.num_layers):
            h, c = hidden_state[layer_idx]
            output_inner = []
            for t in range(seq_len):
                h, c = self.cell_list[layer_idx](input_tensor=cur_layer_input[:, t, :, :, :], h=h, c=c)
                output_inner.append(h)

            layer_output = torch.stack(output_inner, dim=1)
            cur_layer_input = layer_output

            layer_output_list.append(layer_output)
            last_state_list.append((h, c))

        if self.batch_first:
            return layer_output_list, last_state_list
        else:
            return layer_output_list.transpose(0, 1), last_state_list

    def _init_hidden(self, batch_size):
        init_states = []
        for i in range(self.num_layers):
            init_states.append((Variable(torch.zeros(batch_size, self.hidden_size, 64, 64)).cuda(),
                                Variable(torch.zeros(batch_size, self.hidden_size, 64, 64)).cuda()))
        return init_states

这是一个ConvLSTM的实现，可以直接调用。