text: import numpy as np
  import torch
  import torch.nn as nn
  
  from network_tools import scale_tensor
  
  """
  The MS-Net 
  A Computationally Efficient Multiscale Neural Network 
  
  
  The code to generate each individual conv model was modified from:
  https://github.com/tamarott/SinGAN
  
  """
  
  
  def get_trainable_models(scales, features, filters, f_mult):
      
      """
      Returns an array with n-trainable models (ConvNets)
      """
      
      models   = []         # empty list to store the models
      nc_in    = features   # number of inputs on the first layer
      norm     = True       # use Norm
      last_act = None       # activation function
      
      # list of number filters in each model (scale)
      num_filters = [ filters*f_mult**scale for scale in range(scales) ][::-1]
      print(f'Filters per model: {num_filters}')
     
      for it in range( scales ): # creates a model for each scale
          if it==1: nc_in+=1     # adds an additional input to the subsecuent models 
                                 # to convolve the domain + previous(upscaled) result 
          models.append( get_model( nc_in    = nc_in,
                                    ncf      = num_filters[it],
                                    norm     = norm,
                                    last_act = last_act) )
      return models  
  
  
  class get_model(nn.Module):
      def __init__(self, nc_in, ncf, norm, last_act):
          super(get_model, self).__init__()
          
          # default parameters
          nc_out     = 1   # number of output channels of the last layer
          ker_size   = 3   # kernel side-lenght
          padd_size  = 1   # padding size
          ncf_min    = ncf # min number of convolutional filters
          num_layers = 5   # number of conv layers
          
          # first block
          self.head = ConvBlock3D( in_channel  = nc_in,
                                   out_channel = ncf,
                                   ker_size    = ker_size,
                                   padd        = padd_size,
                                   stride      = 1,
                                   norm        = norm )
          
          # body of the model
          self.body = nn.Sequential()
          for i in range( num_layers-1 ):
              new_ncf = int( ncf/2**(i+1) )
              if i==num_layers-2: norm=False  # no norm in the penultimate block
            
              convblock = ConvBlock3D( in_channel  = max(2*new_ncf,ncf_min),
                                       out_channel = max(new_ncf,ncf_min),
                                       ker_size    = ker_size,
                                       padd        = padd_size,
                                       stride      = 1,
                                       norm        = norm)
              
              self.body.add_module( f'block{i+1}', convblock )
          
          if last_act == 'CELU':
              self.tail = nn.Sequential(
                                      nn.Conv3d( max(new_ncf,ncf_min), nc_out,
                                                 kernel_size=1,stride=1, padding=0),
                                      nn.CELU()
                                   )
          else:
              self.tail = nn.Sequential(
                              nn.Conv3d( max(new_ncf,ncf_min), nc_out, kernel_size=1,
                                         stride=1, padding=0)) # no pad needed since 1x1x1
              
      def forward(self,x):
          x = self.head(x)
          x = self.body(x)
          x = self.tail(x)
          return x
  
  
  
  
  class ConvBlock3D( nn.Sequential ):
      def __init__(self, in_channel, out_channel, ker_size, padd, stride, norm):
          super(ConvBlock3D,self).__init__()
          self.add_module( 'conv',
                           nn.Conv3d( in_channel, 
                                      out_channel,
                                      kernel_size=ker_size,
                                      stride=stride,
                                      padding=padd ) ),
          if norm == True:
              self.add_module( 'i_norm', nn.InstanceNorm3d( out_channel ) ),
          self.add_module( 'CeLU', nn.CELU( inplace=False ) )
          
          
  
  class MS_Net(nn.Module):
      
      def __init__(
                   self, 
                   net_name     = 'test1', 
                   num_scales   =  4,
                   num_features =  1, 
                   num_filters  =  2, 
                   f_mult       =  4,  
                   device       =  'cpu',
                   
                   summary      = False
                   ):
          
          super(MS_Net, self).__init__()
          
          self.net_name = net_name
          self.scales   = num_scales
          self.feats    = num_features
          self.feats    = num_features
          self.device   = device
          
          self.models   = nn.ModuleList( 
                                  get_trainable_models( num_scales,
                                                        num_features,
                                                        num_filters,
                                                        f_mult ) )
          if summary:
              print(f'\n Here is a summary of your MS-Net ({net_name}): \n {self.models}')
          
          
          
          
      def forward(self, x_list, masks):
          assert x_list[0].shape[1] == self.feats, \
          f'The number of features provided {x_list[0].shape[1]} \
              does not match with the input size {self.feats}'
          # Carry-out the first prediction (pass through the coarsest model)
          y = [ self.models[0]( x_list[0] ) ]
          for scale,[ model,x ] in enumerate(zip( self.models[1:],x_list[1:] )):
              y_up = scale_tensor( y[scale], scale_factor=2 )*masks[scale]
              y.append( model( torch.cat((x,y_up),dim=1) ) + y_up )
          return y
          
          
      def scale_init(self):
          factor  = 0.3*np.e**(1.9*self.scales)/2
          for num, model in enumerate( self.models ):
              model.tail[0].weight = torch.nn.Parameter(model.tail[0].weight/factor)
              #model.tail[0].bias   = torch.nn.Parameter(model.tail[0].bias/factor)
              
              
      def test(self, masks, x1, x2):
          with torch.no_grad():
              y_up    = [0]*self.scales # upscaled prediction from coarser model
              x = [ torch.cat((xi1, xi2), dim=1) for xi1, xi2 in zip(x1,x2) ]
              del x1, x2
              for scale, model in enumerate( self.models ):
                  if scale == 0:
                      y_up[scale] = model( x[scale] ) 
                  else:
                      y_up = [scale_tensor(y,scale_factor=2)*masks[scale-1] if torch.is_tensor(y) else 0 for y in y_up]
                      y_scale = sum(y_up)
                      y_up[scale] = model(torch.cat((x[scale],y_scale), dim=1))+y_scale
                      del y_scale
              return y_up[-1]
      
      def save_model(self):
          torch.save(self.models, 
                     f'savedModels/{self.net_name}/{self.net_name}.pt')
          
      def save_as(self, name):
          torch.save(self.models, 
                     f'savedModels/{self.net_name}/{self.net_name}_{name}.pt')
          
      def load_model(self):
          try:
              self.models=torch.load(
                  f'savedModels/{self.net_name}/{self.net_name}.pt',
                  map_location=self.device)
              print('Model loaded succesfully')
              return True
          except FileNotFoundError:
              print('No pre-trained model found')
              return False
  
      def load_as(self, name):
          try:
              self.models=torch.load(
                  f'savedModels/{self.net_name}/{self.net_name}_{name}.pt',
                  map_location=self.device)
              print('Model loaded succesfully')
              return True
          except FileNotFoundError:
              print(f'No model called {name} found')
              return False
          
      def num_params(self):
          return sum(p.numel() for p in self.models.parameters() if p.requires_grad)
      
      def maxsize_singlescale(self):
          with torch.set_grad_enabled( False ):
              for cube_size in np.arange(256,1000,64):
                  print(f'Trying out size {cube_size}')
                  try:
                      y = self.models[-1](torch.ones([1,3,cube_size,cube_size,cube_size]).to('cuda'))
                      del y
                  except RuntimeError:
                      return cube_size-64