code: """
  The class wrapper for the networks
  """
  # Built-in
  import os
  import time
  
  # Torch
  import torch
  from torch import nn
  from torch.utils.tensorboard import SummaryWriter
  from torchsummary import summary
  
  # Libs
  import numpy as np
  
  # Own module
  
  
  class Network(object):
      def __init__(self, model_fn, flags, train_loader, test_loader,
                   ckpt_dir=os.path.join(os.path.abspath(''), 'models'),
                   inference_mode=False, saved_model=None):
          self.model_fn = model_fn                                # The model maker function
          self.flags = flags                                      # The Flags containing the specs
          if inference_mode:                                      # If inference mode, use saved model
              self.ckpt_dir = os.path.join(ckpt_dir, saved_model)
              self.saved_model = saved_model
          else:                                                   # training mode, create a new ckpt folder
              self.ckpt_dir = os.path.join(ckpt_dir, time.strftime('%Y%m%d_%H%M%S', time.localtime()))
          self.model = self.create_model()                        # The model itself
          self.loss = self.make_loss()                            # The loss function
          self.optm = self.make_optimizer()                       # The optimizer
          self.train_loader = train_loader                        # The train data loader
          self.test_loader = test_loader                          # The test data loader
          self.log = SummaryWriter(self.ckpt_dir)     # Create a summary writer for keeping the summary to the tensor board
          self.best_validation_loss = float('inf')    # Set the BVL to large number
  
      def create_model(self):
          """
          Function to create the network module from provided model fn and flags
          :return: the created nn module
          """
          model = self.model_fn(self.flags)
          #summary(model, input_size=(128, 8))
          print(model)
          return model
  
      def make_loss(self, logit=None, labels=None):
          """
          Create a tensor that represents the loss. This is consistant both at training time \
          and inference time for Backward model
          :param logit: The output of the network
          :return: the total loss
          """
          if logit is None:
              return None
          MSE_loss = nn.functional.mse_loss(logit, labels)          # The MSE Loss of the
          BDY_loss = 0 # Implemenation later
          return MSE_loss + BDY_loss
  
      def make_optimizer(self):
          """
          Make the corresponding optimizer from the flags. Only below optimizers are allowed. Welcome to add more
          :return:
          """
          if self.flags.optim == 'Adam':
              op = torch.optim.Adam(self.model.parameters(), lr=self.flags.lr, weight_decay=self.flags.reg_scale)
          elif self.flags.optim == 'RMSprop':
              op = torch.optim.RMSprop(self.model.parameters(), lr=self.flags.lr, weight_decay=self.flags.reg_scale)
          elif self.flags.optim == 'SGD':
              op = torch.optim.SGD(self.model.parameters(), lr=self.flags.lr, weight_decay=self.flags.reg_scale)
          else:
              raise Exception("Your Optimizer is neither Adam, RMSprop or SGD, please change in param or contact Ben")
          return op
  
      def save(self):
          """
          Saving the model to the current check point folder with name best_model.pt
          :return: None
          """
          #torch.save(self.model.state_dict, os.path.join(self.ckpt_dir, 'best_model_state_dict.pt'))
          torch.save(self.model, os.path.join(self.ckpt_dir, 'best_model.pt'))
  
      def load(self):
          """
          Loading the model from the check point folder with name best_model.pt
          :return:
          """
          #self.model.load_state_dict(torch.load(os.path.join(self.ckpt_dir, 'best_model_state_dict.pt')))
          self.model.load(torch.load(os.path.join(self.ckpt_dir, 'best_model.pt')))
  
      def train(self):
          """
          The major training function. This would start the training using information given in the flags
          :return: None
          """
          cuda = True if torch.cuda.is_available() else False
          if cuda:
              self.model.cuda()
          for epoch in range(self.flags.train_step):
              # Set to Training Mode
              train_loss = 0
              self.model.train()
              for j, (geometry, spectra) in enumerate(self.train_loader):
                  if cuda:
                      geometry = geometry.cuda()                          # Put data onto GPU
                      spectra = spectra.cuda()                            # Put data onto GPU
                  self.optm.zero_grad()                               # Zero the gradient first
                  logit = self.model(geometry)                        # Get the output
                  loss = self.make_loss(logit, spectra)              # Get the loss tensor
                  loss.backward()                                # Calculate the backward gradients
                  self.optm.step()                                    # Move one step the optimizer
                  train_loss += loss                                  # Aggregate the loss
  
              if epoch % self.flags.eval_step:                        # For eval steps, do the evaluations and tensor board
                  # Record the training loss to the tensorboard
                  train_avg_loss = train_loss.data.numpy() / (j+1)
                  self.log.add_scalar('Loss/train', train_avg_loss, epoch)
  
                  # Set to Evaluation Mode
                  self.model.eval()
                  print("Doing Evaluation on the model now")
                  test_loss = 0
                  for j, (geometry, spectra) in enumerate(self.test_loader):  # Loop through the eval set
                      if cuda:
                          geometry = geometry.cuda()
                          spectra = spectra.cuda()
                      logit = self.model(geometry)
                      loss = self.make_loss(logit, spectra)                   # compute the loss
                      test_loss += loss                                       # Aggregate the loss
  
                  # Record the testing loss to the tensorboard
                  test_avg_loss = test_loss.data.numpy() / (j+1)
                  self.log.add_scalar('Loss/test', test_avg_loss, epoch)
  
                  print("This is Epoch %d, training loss %.5f, validation loss %.5f" \
                        % (epoch, train_avg_loss, test_avg_loss ))
  
                  # Model improving, save the model down
                  if test_avg_loss < self.best_validation_loss:
                      self.best_validation_loss = test_avg_loss
                      self.save()
                      print("Saving the model down...")
  
                      if self.best_validation_loss < self.flags.stop_threshold:
                          print("Training finished EARLIER at epoch %d, reaching loss of %.5f" %\
                                (epoch, self.best_validation_loss))
                          return None
  
      def evaluate(self, save_dir='data/'):
          self.load()
          self.model.eval()                       # Evaluation mode
  
          # Get the file names
          Ypred_file = os.path.join(save_dir, 'test_Ypred_{}.csv'.format(self.saved_model))
          Xtruth_file = os.path.join(save_dir, 'test_Xtruth_{}.csv'.format(self.saved_model))
          Ytruth_file = os.path.join(save_dir, 'test_Ytruth_{}.csv'.format(self.saved_model))
          #Xpred_file = os.path.join(save_dir, 'test_Xpred_{}.csv'.format(self.saved_model))  # For pure forward model, there is no Xpred
  
          # Open those files to append
          with open(Xtruth_file,'a') as fxt,open(Ytruth_file, 'a') as fyt, open(Ypred_file,'a') as fyp:
              # Loop through the eval data and evaluate
              for ind, (geometry, spectra) in enumerate(self.test_loader):
                  logits = self.model(geometry)
                  np.savetxt(fxt, geometry.numpy(), fmt='%.3f')
                  np.savetxt(fyt, spectra.numpy(), fmt='%.3f')
                  np.savetxt(fyp, logits.numpy(), fmt='%.3f')
  
  
  
apis: torch.utils.tensorboard.SummaryWriter torch.nn.functional.mse_loss torch.cuda.is_available