remove labml_helpers dependency: replace Module with nn.Module

Author: vpj

labml_nn/activations/fta/experiment.py

@@ -25,7 +25,6 @@
 
 from labml import experiment
 from labml.configs import option
-from labml_helpers.module import Module
 from labml_nn.activations.fta import FTA
 from labml_nn.experiments.nlp_autoregression import NLPAutoRegressionConfigs
 from labml_nn.transformers import MultiHeadAttention, TransformerLayer
@@ -65,7 +64,7 @@ def forward(self, x: torch.Tensor):
         return self.layer2(x)
 
 
-class AutoregressiveTransformer(Module):
+class AutoregressiveTransformer(nn.Module):
     """
     ## Auto-Regressive model
 

labml_nn/activations/swish.py

@@ -1,10 +1,9 @@
 import torch
 from torch import nn
 
-from labml_helpers.module import Module
 
 
-class Swish(Module):
+class Swish(nn.Module):
     def __init__(self):
         super().__init__()
         self.sigmoid = nn.Sigmoid()

labml_nn/adaptive_computation/ponder_net/__init__.py

@@ -65,10 +65,9 @@
 import torch
 from torch import nn
 
-from labml_helpers.module import Module
 
 
-class ParityPonderGRU(Module):
+class ParityPonderGRU(nn.Module):
     """
     ## PonderNet with GRU for Parity Task
 
@@ -177,7 +176,7 @@ def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Te
         return torch.stack(p), torch.stack(y), p_m, y_m
 
 
-class ReconstructionLoss(Module):
+class ReconstructionLoss(nn.Module):
     """
     ## Reconstruction loss
 
@@ -213,7 +212,7 @@ def forward(self, p: torch.Tensor, y_hat: torch.Tensor, y: torch.Tensor):
         return total_loss
 
 
-class RegularizationLoss(Module):
+class RegularizationLoss(nn.Module):
     """
     ## Regularization loss
 

labml_nn/capsule_networks/__init__.py

@@ -33,10 +33,8 @@
 import torch.nn.functional as F
 import torch.utils.data
 
-from labml_helpers.module import Module
 
-
-class Squash(Module):
+class Squash(nn.Module):
     """
     ## Squash
 
@@ -70,7 +68,7 @@ def forward(self, s: torch.Tensor):
         return (s2 / (1 + s2)) * (s / torch.sqrt(s2 + self.epsilon))
 
 
-class Router(Module):
+class Router(nn.Module):
     """
     ## Routing Algorithm
 
@@ -133,7 +131,7 @@ def forward(self, u: torch.Tensor):
         return v
 
 
-class MarginLoss(Module):
+class MarginLoss(nn.Module):
     """
     ## Margin loss for class existence
 
@@ -153,6 +151,7 @@ class MarginLoss(Module):
     The $\lambda$ down-weighting is used to stop the length of all capsules from
     falling during the initial phase of training.
     """
+
     def __init__(self, *, n_labels: int, lambda_: float = 0.5, m_positive: float = 0.9, m_negative: float = 0.1):
         super().__init__()
 

labml_nn/capsule_networks/mnist.py

@@ -21,12 +21,11 @@
 from labml.configs import option
 from labml_helpers.datasets.mnist import MNISTConfigs
 from labml_helpers.metrics.accuracy import AccuracyDirect
-from labml_helpers.module import Module
 from labml_helpers.train_valid import SimpleTrainValidConfigs, BatchIndex
 from labml_nn.capsule_networks import Squash, Router, MarginLoss
 
 
-class MNISTCapsuleNetworkModel(Module):
+class MNISTCapsuleNetworkModel(nn.Module):
     """
     ## Model for classifying MNIST digits
     """

labml_nn/conv_mixer/__init__.py

@@ -36,11 +36,10 @@
 import torch
 from torch import nn
 
-from labml_helpers.module import Module
 from labml_nn.utils import clone_module_list
 
 
-class ConvMixerLayer(Module):
+class ConvMixerLayer(nn.Module):
     """
     <a id="ConvMixerLayer"></a>
 
@@ -96,7 +95,7 @@ def forward(self, x: torch.Tensor):
         return x
 
 
-class PatchEmbeddings(Module):
+class PatchEmbeddings(nn.Module):
     """
     <a id="PatchEmbeddings"></a>
 
@@ -136,7 +135,7 @@ def forward(self, x: torch.Tensor):
         return x
 
 
-class ClassificationHead(Module):
+class ClassificationHead(nn.Module):
     """
     <a id="ClassificationHead"></a>
 
@@ -169,7 +168,7 @@ def forward(self, x: torch.Tensor):
         return x
 
 
-class ConvMixer(Module):
+class ConvMixer(nn.Module):
     """
     ## ConvMixer
 

labml_nn/diffusion/ddpm/unet.py

@@ -27,10 +27,8 @@
 import torch
 from torch import nn
 
-from labml_helpers.module import Module
 
-
-class Swish(Module):
+class Swish(nn.Module):
     """
     ### Swish activation function
 
@@ -83,7 +81,7 @@ def forward(self, t: torch.Tensor):
         return emb
 
 
-class ResidualBlock(Module):
+class ResidualBlock(nn.Module):
     """
     ### Residual block
 
@@ -140,7 +138,7 @@ def forward(self, x: torch.Tensor, t: torch.Tensor):
         return h + self.shortcut(x)
 
 
-class AttentionBlock(Module):
+class AttentionBlock(nn.Module):
     """
     ### Attention block
 
@@ -208,7 +206,7 @@ def forward(self, x: torch.Tensor, t: Optional[torch.Tensor] = None):
         return res
 
 
-class DownBlock(Module):
+class DownBlock(nn.Module):
     """
     ### Down block
 
@@ -229,7 +227,7 @@ def forward(self, x: torch.Tensor, t: torch.Tensor):
         return x
 
 
-class UpBlock(Module):
+class UpBlock(nn.Module):
     """
     ### Up block
 
@@ -252,7 +250,7 @@ def forward(self, x: torch.Tensor, t: torch.Tensor):
         return x
 
 
-class MiddleBlock(Module):
+class MiddleBlock(nn.Module):
     """
     ### Middle block
 
@@ -305,7 +303,7 @@ def forward(self, x: torch.Tensor, t: torch.Tensor):
         return self.conv(x)
 
 
-class UNet(Module):
+class UNet(nn.Module):
     """
     ## U-Net
     """

labml_nn/experiments/cifar10.py

@@ -14,7 +14,6 @@
 from labml import lab
 from labml.configs import option
 from labml_helpers.datasets.cifar10 import CIFAR10Configs as CIFAR10DatasetConfigs
-from labml_helpers.module import Module
 from labml_nn.experiments.mnist import MNISTConfigs
 
 
@@ -67,7 +66,7 @@ def cifar10_valid_no_augment():
                    ]))
 
 
-class CIFAR10VGGModel(Module):
+class CIFAR10VGGModel(nn.Module):
     """
     ### VGG model for CIFAR-10 classification
     """

labml_nn/experiments/mnist.py

@@ -10,7 +10,6 @@
 
 import torch.nn as nn
 import torch.utils.data
-from labml_helpers.module import Module
 
 from labml import tracker
 from labml.configs import option
@@ -34,7 +33,7 @@ class MNISTConfigs(MNISTDatasetConfigs, TrainValidConfigs):
     device: torch.device = DeviceConfigs()
 
     # Classification model
-    model: Module
+    model: nn.Module
     # Number of epochs to train for
     epochs: int = 10
 

labml_nn/experiments/nlp_autoregression.py

@@ -20,12 +20,11 @@
 from labml_helpers.datasets.text import TextDataset, SequentialDataLoader, SequentialUnBatchedDataset, TextFileDataset
 from labml_helpers.device import DeviceConfigs
 from labml_helpers.metrics.accuracy import Accuracy
-from labml_helpers.module import Module
 from labml_helpers.train_valid import TrainValidConfigs, hook_model_outputs, BatchIndex
 from labml_nn.optimizers.configs import OptimizerConfigs
 
 
-class CrossEntropyLoss(Module):
+class CrossEntropyLoss(nn.Module):
     """
     ### Cross entropy loss
     """
@@ -54,7 +53,7 @@ class NLPAutoRegressionConfigs(TrainValidConfigs):
     device: torch.device = DeviceConfigs()
 
     # Autoregressive model
-    model: Module
+    model: nn.Module
     # Text dataset
     text: TextDataset
     # Batch size