cleanup

Author: vpj

docs/transformers/basic/autoregressive_experiment.html

@@ -76,10 +76,10 @@ <h1>Transformer Auto-Regression Experiment</h1>
         </div>
         <div class='code'>
             <div class="highlight"><pre><span class="lineno">17</span><span></span><span class="kn">import</span> <span class="nn">torch</span>
-<span class="lineno">18</span>
-<span class="lineno">19</span><span class="kn">from</span> <span class="nn">labml</span> <span class="kn">import</span> <span class="n">experiment</span>
-<span class="lineno">20</span><span class="kn">from</span> <span class="nn">labml.configs</span> <span class="kn">import</span> <span class="n">option</span>
-<span class="lineno">21</span><span class="kn">from</span> <span class="nn">labml_helpers.module</span> <span class="kn">import</span> <span class="n">Module</span>
+<span class="lineno">18</span><span class="kn">from</span> <span class="nn">torch</span> <span class="kn">import</span> <span class="n">nn</span>
+<span class="lineno">19</span>
+<span class="lineno">20</span><span class="kn">from</span> <span class="nn">labml</span> <span class="kn">import</span> <span class="n">experiment</span>
+<span class="lineno">21</span><span class="kn">from</span> <span class="nn">labml.configs</span> <span class="kn">import</span> <span class="n">option</span>
 <span class="lineno">22</span><span class="kn">from</span> <span class="nn">labml_nn.experiments.nlp_autoregression</span> <span class="kn">import</span> <span class="n">NLPAutoRegressionConfigs</span>
 <span class="lineno">23</span><span class="kn">from</span> <span class="nn">labml_nn.transformers</span> <span class="kn">import</span> <span class="n">TransformerConfigs</span><span class="p">,</span> <span class="n">Encoder</span>
 <span class="lineno">24</span><span class="kn">from</span> <span class="nn">labml_nn.transformers.utils</span> <span class="kn">import</span> <span class="n">subsequent_mask</span></pre></div>
@@ -94,7 +94,7 @@ <h2>Auto-Regressive model</h2>
 
         </div>
         <div class='code'>
-            <div class="highlight"><pre><span class="lineno">27</span><span class="k">class</span> <span class="nc">AutoregressiveTransformer</span><span class="p">(</span><span class="n">Module</span><span class="p">):</span></pre></div>
+            <div class="highlight"><pre><span class="lineno">27</span><span class="k">class</span> <span class="nc">AutoregressiveTransformer</span><span class="p">(</span><span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">):</span></pre></div>
         </div>
     </div>
     <div class='section' id='section-2'>
@@ -111,7 +111,7 @@ <h2>Auto-Regressive model</h2>
 
         </div>
         <div class='code'>
-            <div class="highlight"><pre><span class="lineno">31</span>    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">encoder</span><span class="p">:</span> <span class="n">Encoder</span><span class="p">,</span> <span class="n">src_embed</span><span class="p">:</span> <span class="n">Module</span><span class="p">,</span> <span class="n">generator</span><span class="p">:</span> <span class="n">Module</span><span class="p">):</span></pre></div>
+            <div class="highlight"><pre><span class="lineno">31</span>    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">encoder</span><span class="p">:</span> <span class="n">Encoder</span><span class="p">,</span> <span class="n">src_embed</span><span class="p">:</span> <span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">,</span> <span class="n">generator</span><span class="p">:</span> <span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">):</span></pre></div>
         </div>
     </div>
     <div class='section' id='section-3'>

docs/transformers/mha.html

@@ -15,20 +15,20 @@
 """
 
 import torch
+from torch import nn
 
 from labml import experiment
 from labml.configs import option
-from labml_helpers.module import Module
 from labml_nn.experiments.nlp_autoregression import NLPAutoRegressionConfigs
 from labml_nn.transformers import TransformerConfigs, Encoder
 from labml_nn.transformers.utils import subsequent_mask
 
 
-class AutoregressiveTransformer(Module):
+class AutoregressiveTransformer(nn.Module):
     """
     ## Auto-Regressive model
     """
-    def __init__(self, encoder: Encoder, src_embed: Module, generator: Module):
+    def __init__(self, encoder: Encoder, src_embed: nn.Module, generator: nn.Module):
         """
         * `encoder` is the transformer [Encoder](../models.html#Encoder)
         * `src_embed` is the token

docs/transformers/models.html

@@ -26,13 +26,12 @@
 from typing import Optional, List
 
 import torch
-from torch import nn as nn
+from torch import nn
 
 from labml import tracker
-from labml_helpers.module import Module
 
 
-class PrepareForMultiHeadAttention(Module):
+class PrepareForMultiHeadAttention(nn.Module):
     """
     <a id="PrepareMHA"></a>
 
@@ -68,7 +67,7 @@ def forward(self, x: torch.Tensor):
         return x
 
 
-class MultiHeadAttention(Module):
+class MultiHeadAttention(nn.Module):
     r"""
     <a id="MHA"></a>
 

docs/transformers/positional_encoding.html

@@ -15,15 +15,14 @@
 
 import torch
 import torch.nn as nn
-from labml_helpers.module import Module
 
 from labml_nn.utils import clone_module_list
 from .feed_forward import FeedForward
 from .mha import MultiHeadAttention
 from .positional_encoding import get_positional_encoding
 
 
-class EmbeddingsWithPositionalEncoding(Module):
+class EmbeddingsWithPositionalEncoding(nn.Module):
     """
     <a id="EmbeddingsWithPositionalEncoding"></a>
 
@@ -41,7 +40,7 @@ def forward(self, x: torch.Tensor):
         return self.linear(x) * math.sqrt(self.d_model) + pe
 
 
-class EmbeddingsWithLearnedPositionalEncoding(Module):
+class EmbeddingsWithLearnedPositionalEncoding(nn.Module):
     """
     <a id="EmbeddingsWithLearnedPositionalEncoding"></a>
 
@@ -59,7 +58,7 @@ def forward(self, x: torch.Tensor):
         return self.linear(x) * math.sqrt(self.d_model) + pe
 
 
-class TransformerLayer(Module):
+class TransformerLayer(nn.Module):
     """
     <a id="TransformerLayer"></a>
 
@@ -139,7 +138,7 @@ def forward(self, *,
         return x
 
 
-class Encoder(Module):
+class Encoder(nn.Module):
     """
     <a id="Encoder"></a>
 
@@ -161,7 +160,7 @@ def forward(self, x: torch.Tensor, mask: torch.Tensor):
         return self.norm(x)
 
 
-class Decoder(Module):
+class Decoder(nn.Module):
     """
     <a id="Decoder"></a>
 
@@ -183,7 +182,7 @@ def forward(self, x: torch.Tensor, memory: torch.Tensor, src_mask: torch.Tensor,
         return self.norm(x)
 
 
-class Generator(Module):
+class Generator(nn.Module):
     """
     <a id="Generator"></a>
 
@@ -201,14 +200,14 @@ def forward(self, x):
         return self.projection(x)
 
 
-class EncoderDecoder(Module):
+class EncoderDecoder(nn.Module):
     """
     <a id="EncoderDecoder"></a>
 
     ## Combined Encoder-Decoder
     """
 
-    def __init__(self, encoder: Encoder, decoder: Decoder, src_embed: Module, tgt_embed: Module, generator: Module):
+    def __init__(self, encoder: Encoder, decoder: Decoder, src_embed: nn.Module, tgt_embed: nn.Module, generator: nn.Module):
         super().__init__()
         self.encoder = encoder
         self.decoder = decoder

labml_nn/transformers/basic/autoregressive_experiment.py

@@ -26,10 +26,8 @@
 import torch
 import torch.nn as nn
 
-from labml_helpers.module import Module
 
-
-class PositionalEncoding(Module):
+class PositionalEncoding(nn.Module):
     def __init__(self, d_model: int, dropout_prob: float, max_len: int = 5000):
         super().__init__()
         self.dropout = nn.Dropout(dropout_prob)