Spaces:
Sleeping
Sleeping
File size: 3,634 Bytes
0874d87 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 |
import torch
import torch.nn as nn
class CTCEncoder(nn.Module):
def __init__(self, num_classes, cnn_output_dim=256, rnn_hidden_dim=256, rnn_layers=3):
"""
CTC Encoder with a CNN feature extractor and LSTM for sequence modeling.
Args:
num_classes (int): Number of output classes for the model.
cnn_output_dim (int): Number of output channels from the CNN.
rnn_hidden_dim (int): Hidden size of the LSTM.
rnn_layers (int): Number of layers in the LSTM.
"""
super(CTCEncoder, self).__init__()
# CNN Feature Extractor
self.feature_extractor = nn.Sequential(
nn.Conv2d(1, 32, kernel_size=3, stride=1, padding=1),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2), # Down-sample by 2
nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2), # Down-sample by another 2
nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1),
nn.ReLU(),
nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1),
nn.ReLU(),
nn.AdaptiveAvgPool2d((1, None)) # Ensure output height is 1
)
# Bidirectional LSTM
self.rnn_hidden_dim = rnn_hidden_dim
self.rnn_layers = rnn_layers
self.cnn_output_dim = cnn_output_dim
self.rnn = nn.LSTM(
input_size=cnn_output_dim, # Output channels from CNN
hidden_size=rnn_hidden_dim,
num_layers=rnn_layers,
batch_first=True,
bidirectional=True
)
# Fully connected layer
self.fc = nn.Linear(rnn_hidden_dim * 2, num_classes)
def compute_input_lengths(self, input_lengths):
"""
Adjusts input lengths based on the CNN's down-sampling operations.
Args:
input_lengths (torch.Tensor): Original input lengths.
Returns:
torch.Tensor: Adjusted input lengths.
"""
# Account for down-sampling by MaxPool layers (factor of 2 for each MaxPool)
input_lengths = input_lengths // 2 # First MaxPool
input_lengths = input_lengths // 2 # Second MaxPool
input_lengths = input_lengths // 2 # Third pooling layer or additional down-sampling
return input_lengths
def forward(self, x, input_lengths):
"""
Forward pass through the encoder.
Args:
x (torch.Tensor): Input tensor of shape [B, 1, H, W].
input_lengths (torch.Tensor): Lengths of the sequences in the batch.
Returns:
torch.Tensor: Logits of shape [B, T, num_classes].
torch.Tensor: Adjusted input lengths.
"""
# Feature extraction
x = self.feature_extractor(x) # [Batch_Size, Channels, Height, Width]
print(f"Shape after CNN: {x.shape}") # Debug the shape
# Reshape for LSTM
x = x.squeeze(2).permute(0, 2, 1) # [Batch_Size, Sequence_Length, Features]
assert x.size(-1) == 256, f"Expected last dimension to be 256, but got {x.size(-1)}"
# Adjust input lengths
input_lengths = self.compute_input_lengths(input_lengths)
assert input_lengths.size(0) == x.size(0), f"input_lengths size ({input_lengths.size(0)}) must match batch size ({x.size(0)})"
# Pass through LSTM
x, _ = self.rnn(x) # [Batch_Size, Sequence_Length, 2 * Hidden_Dim]
# Fully connected output
x = self.fc(x) # [Batch_Size, Sequence_Length, Num_Classes]
return x, input_lengths
|