class_docstr: This class tests the crypten.nn package.
class_docstr_tok: ["This","class","tests","the","crypten.nn","package","."]
docstr: Tests training CrypTen models created directly using the crypten.nn.Module
code: def test_custom_module_training(self):
      """Tests training CrypTen models created directly using the crypten.nn.Module"""
      BATCH_SIZE = 32
      NUM_FEATURES = 3
  
      class ExampleNet(crypten.nn.Module):
          def __init__(self):
              super(ExampleNet, self).__init__()
              self.fc1 = crypten.nn.Linear(NUM_FEATURES, BATCH_SIZE)
              self.fc2 = crypten.nn.Linear(BATCH_SIZE, 2)
  
          def forward(self, x):
              out = self.fc1(x)
              out = self.fc2(out)
              return out
  
      model = ExampleNet()
  
      x_orig = get_random_test_tensor(
          size=(BATCH_SIZE, NUM_FEATURES), is_float=True
      )
      # y is a linear combo of x to ensure network can easily learn pattern
      y_orig = (2 * x_orig.mean(dim=1)).gt(0).long()
      y_one_hot = onehot(y_orig, num_targets=2)
  
      # encrypt training sample:
      x_train = crypten.cryptensor(x_orig, requires_grad=True)
      y_train = crypten.cryptensor(y_one_hot)
  
      for loss_name in ["BCELoss", "CrossEntropyLoss", "MSELoss"]:
          # create loss function
          loss = getattr(crypten.nn, loss_name)()
  
          # create encrypted model
          model.train()
          model.encrypt()
  
          num_epochs = 3
          learning_rate = 0.001
  
          for i in range(num_epochs):
              output = model(x_train)
              if loss_name == "MSELoss":
                  output_norm = output
              else:
                  output_norm = output.softmax(1)
              loss_value = loss(output_norm, y_train)
  
              # set gradients to "zero"
              model.zero_grad()
              for param in model.parameters():
                  self.assertIsNone(
                      param.grad, "zero_grad did not reset gradients"
                  )
  
              # perform backward pass:
              loss_value.backward()
              for param in model.parameters():
                  if param.requires_grad:
                      self.assertIsNotNone(
                          param.grad,
                          "required parameter gradient not created",
                      )
  
              # update parameters
              orig_parameters, upd_parameters = {}, {}
              orig_parameters = self._compute_reference_parameters(
                  "", orig_parameters, model, 0
              )
              model.update_parameters(learning_rate)
              upd_parameters = self._compute_reference_parameters(
                  "", upd_parameters, model, learning_rate
              )
  
              parameter_changed = False
              for name, value in orig_parameters.items():
                  if param.requires_grad and param.grad is not None:
                      unchanged = torch.allclose(
                          upd_parameters[name], value
                      )
                      if unchanged is False:
                          parameter_changed = True
                      self.assertTrue(
                          parameter_changed,
                          "no parameter changed in training step",
                      )
  
              # record initial and current loss
              if i == 0:
                  orig_loss = loss_value.get_plain_text()
              curr_loss = loss_value.get_plain_text()
  
          # check that the loss has decreased after training
          self.assertTrue(
              curr_loss.item() < orig_loss.item(),
              "loss has not decreased after training",
          )
  
imports: import copy
  import itertools
  import logging
  import unittest
  import crypten
  import crypten.communicator
  import torch
  import torch.nn.functional
  from crypten.common.rng import generate_random_ring_element
  from crypten.common.tensor_types import is_float_tensor
  from crypten.config import cfg
  from crypten.encoder import FixedPointEncoder
  from test.multiprocess_test_case import get_random_linear, get_random_test_tensor, MultiProcessTestCase, onehot
label_desc: Use image node_id 20 for calling the TestNN obj's underlying member method code with example usage: obj.test_custom_module_training() and returns: out