Spaces:
Sleeping
Sleeping
| 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 | |