fixes

Author: sergeyk

lab2/text_recognizer/data/emnist.py

@@ -14,7 +14,7 @@
 import toml
 import torch
 
-from text_recognizer.data.base_data_module import _download_raw_dataset, BaseDataModule, load_and_print_info
+from text_recognizer.data.base_data_module import _download_raw_dataset, BaseDataModule, load_and_print_info, split_dataset
 from text_recognizer.data.util import BaseDataset
 
 NUM_SPECIAL_TOKENS = 4
@@ -68,11 +68,7 @@ def setup(self, stage: str = None):
                 self.y_trainval = f["y_train"][:].squeeze().astype(int)
 
             data_trainval = BaseDataset(self.x_trainval, self.y_trainval, transform=self.transform)
-            train_size = int(TRAIN_FRAC * len(data_trainval))
-            val_size = len(data_trainval) - train_size
-            self.data_train, self.data_val = torch.utils.data.random_split(
-                data_trainval, [train_size, val_size], generator=torch.Generator().manual_seed(42)
-            )
+            self.data_train, self.data_val = split_dataset(base_dataset=data_trainval, fraction=TRAIN_FRAC, seed=42)
 
         if stage == "test" or stage is None:
             with h5py.File(PROCESSED_DATA_FILENAME, "r") as f:

lab3/text_recognizer/data/emnist.py

@@ -14,7 +14,7 @@
 import toml
 import torch
 
-from text_recognizer.data.base_data_module import _download_raw_dataset, BaseDataModule, load_and_print_info
+from text_recognizer.data.base_data_module import _download_raw_dataset, BaseDataModule, load_and_print_info, split_dataset
 from text_recognizer.data.util import BaseDataset
 
 NUM_SPECIAL_TOKENS = 4
@@ -68,11 +68,7 @@ def setup(self, stage: str = None):
                 self.y_trainval = f["y_train"][:].squeeze().astype(int)
 
             data_trainval = BaseDataset(self.x_trainval, self.y_trainval, transform=self.transform)
-            train_size = int(TRAIN_FRAC * len(data_trainval))
-            val_size = len(data_trainval) - train_size
-            self.data_train, self.data_val = torch.utils.data.random_split(
-                data_trainval, [train_size, val_size], generator=torch.Generator().manual_seed(42)
-            )
+            self.data_train, self.data_val = split_dataset(base_dataset=data_trainval, fraction=TRAIN_FRAC, seed=42)
 
         if stage == "test" or stage is None:
             with h5py.File(PROCESSED_DATA_FILENAME, "r") as f:

lab3/text_recognizer/models/line_cnn.py

@@ -128,6 +128,8 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         x = self.dropout(x)
         x = self.fc2(x)  # (B, S, C)
         x = x.permute(0, 2, 1)  # -> (B, C, S)
+        if self.limit_output_length:
+            x = x[:, :, :self.output_length]
         return x
 
     @staticmethod

lab4/text_recognizer/data/emnist.py

@@ -14,7 +14,7 @@
 import toml
 import torch
 
-from text_recognizer.data.base_data_module import _download_raw_dataset, BaseDataModule, load_and_print_info
+from text_recognizer.data.base_data_module import _download_raw_dataset, BaseDataModule, load_and_print_info, split_dataset
 from text_recognizer.data.util import BaseDataset
 
 NUM_SPECIAL_TOKENS = 4
@@ -68,11 +68,7 @@ def setup(self, stage: str = None):
                 self.y_trainval = f["y_train"][:].squeeze().astype(int)
 
             data_trainval = BaseDataset(self.x_trainval, self.y_trainval, transform=self.transform)
-            train_size = int(TRAIN_FRAC * len(data_trainval))
-            val_size = len(data_trainval) - train_size
-            self.data_train, self.data_val = torch.utils.data.random_split(
-                data_trainval, [train_size, val_size], generator=torch.Generator().manual_seed(42)
-            )
+            self.data_train, self.data_val = split_dataset(base_dataset=data_trainval, fraction=TRAIN_FRAC, seed=42)
 
         if stage == "test" or stage is None:
             with h5py.File(PROCESSED_DATA_FILENAME, "r") as f:

lab4/text_recognizer/models/line_cnn.py

@@ -128,6 +128,8 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         x = self.dropout(x)
         x = self.fc2(x)  # (B, S, C)
         x = x.permute(0, 2, 1)  # -> (B, C, S)
+        if self.limit_output_length:
+            x = x[:, :, :self.output_length]
         return x
 
     @staticmethod

lab5/text_recognizer/data/emnist.py

@@ -14,7 +14,7 @@
 import toml
 import torch
 
-from text_recognizer.data.base_data_module import _download_raw_dataset, BaseDataModule, load_and_print_info
+from text_recognizer.data.base_data_module import _download_raw_dataset, BaseDataModule, load_and_print_info, split_dataset
 from text_recognizer.data.util import BaseDataset
 
 NUM_SPECIAL_TOKENS = 4
@@ -68,11 +68,7 @@ def setup(self, stage: str = None):
                 self.y_trainval = f["y_train"][:].squeeze().astype(int)
 
             data_trainval = BaseDataset(self.x_trainval, self.y_trainval, transform=self.transform)
-            train_size = int(TRAIN_FRAC * len(data_trainval))
-            val_size = len(data_trainval) - train_size
-            self.data_train, self.data_val = torch.utils.data.random_split(
-                data_trainval, [train_size, val_size], generator=torch.Generator().manual_seed(42)
-            )
+            self.data_train, self.data_val = split_dataset(base_dataset=data_trainval, fraction=TRAIN_FRAC, seed=42)
 
         if stage == "test" or stage is None:
             with h5py.File(PROCESSED_DATA_FILENAME, "r") as f:

lab5/text_recognizer/data/iam_lines.py

@@ -16,7 +16,7 @@
 from torchvision import transforms
 
 from text_recognizer.data.util import BaseDataset, convert_strings_to_labels
-from text_recognizer.data.base_data_module import BaseDataModule, load_and_print_info
+from text_recognizer.data.base_data_module import BaseDataModule, load_and_print_info, split_dataset
 from text_recognizer.data.emnist import EMNIST
 from text_recognizer.data.iam import IAM
 from text_recognizer import util
@@ -82,11 +82,7 @@ def setup(self, stage: str = None):
             y_trainval = convert_strings_to_labels(labels_trainval, self.inverse_mapping, length=self.output_dims[0])
             data_trainval = BaseDataset(x_trainval, y_trainval, transform=get_transform(IMAGE_WIDTH, self.augment))
 
-            train_size = int(TRAIN_FRAC * len(data_trainval))
-            val_size = len(data_trainval) - train_size
-            self.data_train, self.data_val = torch.utils.data.random_split(
-                data_trainval, [train_size, val_size], generator=torch.Generator().manual_seed(42)
-            )
+            self.data_train, self.data_val = split_dataset(base_dataset=data_trainval, fraction=TRAIN_FRAC, seed=42)
 
         # Note that test data does not go through augmentation transforms
         if stage == "test" or stage is None:

lab5/text_recognizer/models/line_cnn.py

@@ -128,6 +128,8 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         x = self.dropout(x)
         x = self.fc2(x)  # (B, S, C)
         x = x.permute(0, 2, 1)  # -> (B, C, S)
+        if self.limit_output_length:
+            x = x[:, :, :self.output_length]
         return x
 
     @staticmethod