bilstm_keras_to_pytorch_updated.py

import torch
from torch import nn
class BiLSTM(nn.Module):
    def __init__(self, input_dim, embedding_dim, hidden_dim):
        super().__init__()
        self.input_dim = input_dim
        self.embedding_dim = embedding_dim
        self.hidden_dim = hidden_dim
        self.encoder = nn.Embedding(input_dim, embedding_dim)
        self.lstm = nn.LSTM(embedding_dim, hidden_dim,
                            num_layers=2, bidirectional=True)
        self.linear = nn.Linear(hidden_dim * 2, 1)
        self.activation = nn.Sigmoid()

        nn.init.xavier_uniform_(self.linear.weight)
        self.linear.bias.data.zero_()
        self.init_weights()

    def init_weights(self):
        ih = (param.data for name, param in self.named_parameters() if 'weight_ih' in name)
        hh = (param.data for name, param in self.named_parameters() if 'weight_hh' in name)
        b = (param.data for name, param in self.named_parameters() if 'bias' in name)
        self.encoder.weight.data.uniform_(-0.5, 0.5)
        for t in ih:
            nn.init.xavier_uniform(t)
        for t in hh:
            nn.init.orthogonal(t)
        for t in b:
            nn.init.constant(t, 0)

    def forward(self, src):
        batch_size = src.size(1)
        output = self.encoder(src)
        output, _ = self.lstm(output)
        output = nn.functional.tanh(output[-1])
        output = self.linear(output)
        output = self.activation(output)
        return output, None