class_docstr: ## Model for classifying MNIST digits
class_docstr_tok: ["#","#","Model","for","classifying","MNIST","digits"]
code: def __init__(self):
      super().__init__()
      # First convolution layer has $256$, $9 \times 9$ convolution kernels
      self.conv1 = nn.Conv2d(
          in_channels=1, out_channels=256, kernel_size=9, stride=1
      )
      # The second layer (Primary Capsules) s a convolutional capsule layer with $32$ channels
      # of convolutional $8D$ capsules ($8$ features per capsule).
      # That is, each primary capsule contains 8 convolutional units with a 9 × 9 kernel and a stride of 2.
      # In order to implement this we create a convolutional layer with $32 \times 8$ channels and
      # reshape and permutate its output to get the capsules of $8$ features each.
      self.conv2 = nn.Conv2d(
          in_channels=256,
          out_channels=32 * 8,
          kernel_size=9,
          stride=2,
          padding=0,
      )
      self.squash = Squash()
  
      # Routing layer gets the $32 \times 6 \times 6$ primary capsules and produces $10$ capsules.
      # Each of the primary capsules have $8$ features, while output capsules (Digit Capsules)
      # have $16$ features.
      # The routing algorithm iterates $3$ times.
      self.digit_capsules = Router(32 * 6 * 6, 10, 8, 16, 3)
  
      # This is the decoder mentioned in the paper.
      # It takes the outputs of the $10$ digit capsules, each with $16$ features to reproduce the
      # image. It goes through linear layers of sizes $512$ and $1024$ with $ReLU$ activations.
      self.decoder = nn.Sequential(
          nn.Linear(16 * 10, 512),
          nn.ReLU(),
          nn.Linear(512, 1024),
          nn.ReLU(),
          nn.Linear(1024, 784),
          nn.Sigmoid(),
      )
  
imports: from typing import Any
  import torch.nn
  import torch.nn.functional
  import torch.utils.data
  from labml import experiment, tracker
  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
label_desc: Use image node_id 1 to create a new MNISTCapsuleNetworkModel object from inherited base classes: Module with example: obj = MNISTCapsuleNetworkModel()