如何在 Pytorch 中实现双向 Conv LSTM

如何解决如何在 Pytorch 中实现双向 Conv LSTM

import torch
from torch import nn


def initialize_weights(self,layer):
        """Initialize a layer's weights and biases.

        Args:
            layer: A PyTorch Module's layer."""
        if isinstance(layer,(nn.Batchnorm2d,nn.Batchnorm1d)):
            pass
        else:
            try:
                nn.init.xavier_normal_(layer.weight)
            except AttributeError:
                pass
            try:
                nn.init.uniform_(layer.bias)
            except (ValueError,AttributeError):
                pass
class HadamardProduct(nn.Module):
    """A Hadamard product layer.
    
    Args:
        shape: The shape of the layer."""
       
    def __init__(self,shape):
        super().__init__()
        self.weights = nn.Parameter(torch.empty(*shape))
        self.bias = nn.Parameter(torch.empty(*shape))
           
    def forward(self,x):
        return x * self.weights

    
class ConvLSTMCell(nn.Module):
    """A convolutional LSTM cell.

    Implementation details follow closely the following paper:

    Shi et al. -'Convolutional LSTM Network: A Machine Learning 
    Approach for Precipitation Nowcasting' (2015).
    Accessible at https://arxiv.org/abs/1506.04214

    The parameter names are drawn from the paper's Eq. 3.

    Args:
        input_bands: The number of bands in the input data.
        input_dim: The length of of side of input data. Data is
            presumed to have identical width and heigth."""

    def __init__(self,input_bands,input_dim,kernels,dropout,batch_norm):
        super().__init__()

        self.input_bands = input_bands
        self.input_dim = input_dim
        self.kernels = kernels
        self.dropout = dropout
        self.batch_norm = batch_norm

        self.kernel_size = 3
        self.padding = 1  # Preserve dimensions

        self.input_conv_params = {
            'in_channels': self.input_bands,'out_channels': self.kernels,'kernel_size': self.kernel_size,'padding': self.padding,'bias': True
        }

        self.hidden_conv_params = {
            'in_channels': self.kernels,'bias': True
        }

        self.state_shape = (
            1,self.kernels,self.input_dim,self.input_dim
        )

        self.batch_norm_layer = None
        if self.batch_norm:
            self.batch_norm_layer = nn.Batchnorm2d(num_features=self.input_bands)

        # Input Gates
        self.W_xi = nn.Conv2d(**self.input_conv_params)
        self.W_hi = nn.Conv2d(**self.hidden_conv_params)
        self.W_ci = HadamardProduct(self.state_shape)

        # Forget Gates
        self.W_xf = nn.Conv2d(**self.input_conv_params)
        self.W_hf = nn.Conv2d(**self.hidden_conv_params)
        self.W_cf = HadamardProduct(self.state_shape)

        # Memory Gates
        self.W_xc = nn.Conv2d(**self.input_conv_params)
        self.W_hc = nn.Conv2d(**self.hidden_conv_params)

        # Output Gates
        self.W_xo = nn.Conv2d(**self.input_conv_params)
        self.W_ho = nn.Conv2d(**self.hidden_conv_params)
        self.W_co = HadamardProduct(self.state_shape)

        # Dropouts
        self.H_drop = nn.Dropout2d(p=self.dropout)
        self.C_drop = nn.Dropout2d(p=self.dropout)

        self.apply(initialize_weights)
class ConvLSTM(nn.Module):

    def __init__(self,num_layers,bidirectional,dropout):
        super().__init__()
        self.input_bands = input_bands
        self.input_dim = input_dim
        self.kernels = kernels
        self.num_layers = num_layers
        self.bidirectional = bidirectional
        self.dropout = dropout
        
        self.layers_fwd = self.initialize_layers()
        self.layers_bwd = None
        if self.bidirectional:
            self.layers_bwd = self.initialize_layers()
        self.fc_output = nn.Sequential(
            nn.Flatten(),nn.Linear(
                in_features=self.kernels*self.input_dim**2*(1 if not self.bidirectional else 2),out_features=1024
            ),nn.Linear(
                in_features=1024,out_features=1
            )
        )
            
        self.apply(initialize_weights)
        
    def initialize_layers(self):
        """Initialize a single direction of the model's layers.
        
        This function takes care of stacking layers,allocating
        dropout and assigning correct input feature number for
        each layer in the stack."""
        
        layers = nn.ModuleList()
        
        for i in range(self.num_layers):
            layers.append(
                ConvLSTMCell(
                    input_bands=self.input_bands if i == 0 else self.kernels,input_dim=self.input_dim,dropout=self.dropout if i+1 < self.num_layers else 0,kernels=self.kernels,batch_norm=False
                )
            )
            
        return layers
    
        
    def forward(self,x):
        """Perform forward pass with the model.
        
        For each item in the sequence,the data is propagated 
        through each layer and both directions,if possible.
        In case of a bidirectional model,the outputs are 
        concatenated from both directions. The output of the 
        last item of the sequence is further given to the FC
        layers to produce the final batch of predictions. 
        
        Args:
            x:  A batch of spatial data sequences. The data
                should be in the following format:
                [Batch,Seq,Band,Dim,Dim]
                    
        Returns:
            A batch of predictions."""
        
        seq_len = x.shape[1]
        
        for seq_idx in range(seq_len):
            
            layer_in_out = x[:,seq_idx,::]
            states = None
            for layer in self.layers_fwd:
                layer_in_out,states = layer(layer_in_out,states)
                
            if not self.bidirectional:
                continue
                
            layer_in_out_bwd = x[:,-seq_idx,::]
            states = None
            for layer in self.layers_bwd:
                layer_in_out_bwd,states = layer(layer_in_out_bwd,states)
            
            layer_in_out = torch.cat((layer_in_out,layer_in_out_bwd),dim=1)
            

      return self.fc_output(layer_in_out)

我在 Pytorch 论坛中找到了这个 https://discuss.pytorch.org/t/passing-hidden-layers-to-convlstm/52814 这是我在 Pytorch 中为 Bi-Directional Conv LSTM 找到的实现。在完成代码后，我被困在尝试测试它，但似乎有些函数定义错误或不完整。我在这里 https://github.com/ndrplz/ConvLSTM_pytorch 也找到了 Conv LSTM 的其他实现，但这不支持双向。我需要一些帮助来重新调整上面的代码。

ConvLSTM2D = ConvLSTM(128,128,3,1,True,0.0)
x = torch.randn([5,224,224])
t1 = ConvLSTM2D(x)

print(t1)

TypeError: initialize_weights() missing 1 required positional argument: 'layer'