original_docstring:  Load tf checkpoints and standing statistics in a pytorch model
      
docstring: Load tf checkpoints and standing statistics in a pytorch model
docstring_tokens: Load tf checkpoints and standing statistics in a pytorch model
code: def load_tf_weights_in_biggan(model, config, tf_model_path, batch_norm_stats_path=None):
      """ Load tf checkpoints and standing statistics in a pytorch model
      """
      try:
          import numpy as np
          import tensorflow as tf
      except ImportError:
          raise ImportError("Loading a TensorFlow models in PyTorch, requires TensorFlow to be installed. Please see "
              "https://www.tensorflow.org/install/ for installation instructions.")
      # Load weights from TF model
      checkpoint_path = tf_model_path + "/variables/variables"
      init_vars = tf.train.list_variables(checkpoint_path)
      from pprint import pprint
      pprint(init_vars)
  
      # Extract batch norm statistics from model if needed
      if batch_norm_stats_path:
          stats = torch.load(batch_norm_stats_path)
      else:
          logger.info("Extracting batch norm stats")
          stats = extract_batch_norm_stats(tf_model_path)
  
      # Build TF to PyTorch weights loading map
      tf_to_pt_map = build_tf_to_pytorch_map(model, config)
  
      tf_weights = {}
      for name in tf_to_pt_map.keys():
          array = tf.train.load_variable(checkpoint_path, name)
          tf_weights[name] = array
          # logger.info("Loading TF weight {} with shape {}".format(name, array.shape))
  
      # Load parameters
      with torch.no_grad():
          pt_params_pnt = set()
          for name, pointer in tf_to_pt_map.items():
              array = tf_weights[name]
              if pointer.dim() == 1:
                  if pointer.dim() < array.ndim:
                      array = np.squeeze(array)
              elif pointer.dim() == 2:  # Weights
                  array = np.transpose(array)
              elif pointer.dim() == 4:  # Convolutions
                  array = np.transpose(array, (3, 2, 0, 1))
              else:
                  raise "Wrong dimensions to adjust: " + str((pointer.shape, array.shape))
              if pointer.shape != array.shape:
                  raise ValueError("Wrong dimensions: " + str((pointer.shape, array.shape)))
              logger.info("Initialize PyTorch weight {} with shape {}".format(name, pointer.shape))
              pointer.data = torch.from_numpy(array) if isinstance(array, np.ndarray) else torch.tensor(array)
              tf_weights.pop(name, None)
              pt_params_pnt.add(pointer.data_ptr())
  
          # Prepare SpectralNorm buffers by running one step of Spectral Norm (no need to train the model):
          for module in model.modules():
              for n, buffer in module.named_buffers():
                  if n == 'weight_v':
                      weight_mat = module.weight_orig
                      weight_mat = weight_mat.reshape(weight_mat.size(0), -1)
                      u = module.weight_u
  
                      v = normalize(torch.mv(weight_mat.t(), u), dim=0, eps=config.eps)
                      buffer.data = v
                      pt_params_pnt.add(buffer.data_ptr())
  
                      u = normalize(torch.mv(weight_mat, v), dim=0, eps=config.eps)
                      module.weight_u.data = u
                      pt_params_pnt.add(module.weight_u.data_ptr())
  
          # Load batch norm statistics
          index = 0
          for layer in model.generator.layers:
              if not hasattr(layer, 'bn_0'):
                  continue
              for i in range(4):  #  Batchnorms
                  bn_pointer = getattr(layer, 'bn_%d' % i)
                  pointer = bn_pointer.running_means
                  if pointer.shape != stats[index].shape:
                      raise "Wrong dimensions: " + str((pointer.shape, stats[index].shape))
                  pointer.data = torch.from_numpy(stats[index])
                  pt_params_pnt.add(pointer.data_ptr())
  
                  pointer = bn_pointer.running_vars
                  if pointer.shape != stats[index+1].shape:
                      raise "Wrong dimensions: " + str((pointer.shape, stats[index].shape))
                  pointer.data = torch.from_numpy(stats[index+1])
                  pt_params_pnt.add(pointer.data_ptr())
  
                  index += 2
  
          bn_pointer = model.generator.bn
          pointer = bn_pointer.running_means
          if pointer.shape != stats[index].shape:
              raise "Wrong dimensions: " + str((pointer.shape, stats[index].shape))
          pointer.data = torch.from_numpy(stats[index])
          pt_params_pnt.add(pointer.data_ptr())
  
          pointer = bn_pointer.running_vars
          if pointer.shape != stats[index+1].shape:
              raise "Wrong dimensions: " + str((pointer.shape, stats[index].shape))
          pointer.data = torch.from_numpy(stats[index+1])
          pt_params_pnt.add(pointer.data_ptr())
  
      remaining_params = list(n for n, t in chain(model.named_parameters(), model.named_buffers()) \
                              if t.data_ptr() not in pt_params_pnt)
  
      logger.info("TF Weights not copied to PyTorch model: {} -".format(', '.join(tf_weights.keys())))
      logger.info("Remanining parameters/buffers from PyTorch model: {} -".format(', '.join(remaining_params)))
  
      return model
code_tokens: def load_tf_weights_in_biggan ( model , config , tf_model_path , batch_norm_stats_path = None ) : try : import numpy as np import tensorflow as tf except ImportError : raise ImportError ( "Loading a TensorFlow models in PyTorch, requires TensorFlow to be installed. Please see " "https://www.tensorflow.org/install/ for installation instructions." ) checkpoint_path = tf_model_path + "/variables/variables" init_vars = tf . train . list_variables ( checkpoint_path ) from pprint import pprint pprint ( init_vars ) if batch_norm_stats_path : stats = torch . load ( batch_norm_stats_path ) else : logger . info ( "Extracting batch norm stats" ) stats = extract_batch_norm_stats ( tf_model_path ) tf_to_pt_map = build_tf_to_pytorch_map ( model , config ) tf_weights = { } for name in tf_to_pt_map . keys ( ) : array = tf . train . load_variable ( checkpoint_path , name ) tf_weights [ name ] = array with torch . no_grad ( ) : pt_params_pnt = set ( ) for name , pointer in tf_to_pt_map . items ( ) : array = tf_weights [ name ] if pointer . dim ( ) == 1 : if pointer . dim ( ) < array . ndim : array = np . squeeze ( array ) elif pointer . dim ( ) == 2 : array = np . transpose ( array ) elif pointer . dim ( ) == 4 : array = np . transpose ( array , ( 3 , 2 , 0 , 1 ) ) else : raise "Wrong dimensions to adjust: " + str ( ( pointer . shape , array . shape ) ) if pointer . shape != array . shape : raise ValueError ( "Wrong dimensions: " + str ( ( pointer . shape , array . shape ) ) ) logger . info ( "Initialize PyTorch weight {} with shape {}" . format ( name , pointer . shape ) ) pointer . data = torch . from_numpy ( array ) if isinstance ( array , np . ndarray ) else torch . tensor ( array ) tf_weights . pop ( name , None ) pt_params_pnt . add ( pointer . data_ptr ( ) ) for module in model . modules ( ) : for n , buffer in module . named_buffers ( ) : if n == 'weight_v' : weight_mat = module . weight_orig weight_mat = weight_mat . reshape ( weight_mat . size ( 0 ) , - 1 ) u = module . weight_u v = normalize ( torch . mv ( weight_mat . t ( ) , u ) , dim = 0 , eps = config . eps ) buffer . data = v pt_params_pnt . add ( buffer . data_ptr ( ) ) u = normalize ( torch . mv ( weight_mat , v ) , dim = 0 , eps = config . eps ) module . weight_u . data = u pt_params_pnt . add ( module . weight_u . data_ptr ( ) ) index = 0 for layer in model . generator . layers : if not hasattr ( layer , 'bn_0' ) : continue for i in range ( 4 ) : bn_pointer = getattr ( layer , 'bn_%d' % i ) pointer = bn_pointer . running_means if pointer . shape != stats [ index ] . shape : raise "Wrong dimensions: " + str ( ( pointer . shape , stats [ index ] . shape ) ) pointer . data = torch . from_numpy ( stats [ index ] ) pt_params_pnt . add ( pointer . data_ptr ( ) ) pointer = bn_pointer . running_vars if pointer . shape != stats [ index + 1 ] . shape : raise "Wrong dimensions: " + str ( ( pointer . shape , stats [ index ] . shape ) ) pointer . data = torch . from_numpy ( stats [ index + 1 ] ) pt_params_pnt . add ( pointer . data_ptr ( ) ) index += 2 bn_pointer = model . generator . bn pointer = bn_pointer . running_means if pointer . shape != stats [ index ] . shape : raise "Wrong dimensions: " + str ( ( pointer . shape , stats [ index ] . shape ) ) pointer . data = torch . from_numpy ( stats [ index ] ) pt_params_pnt . add ( pointer . data_ptr ( ) ) pointer = bn_pointer . running_vars if pointer . shape != stats [ index + 1 ] . shape : raise "Wrong dimensions: " + str ( ( pointer . shape , stats [ index ] . shape ) ) pointer . data = torch . from_numpy ( stats [ index + 1 ] ) pt_params_pnt . add ( pointer . data_ptr ( ) ) remaining_params = list ( n for n , t in chain ( model . named_parameters ( ) , model . named_buffers ( ) ) if t . data_ptr ( ) not in pt_params_pnt ) logger . info ( "TF Weights not copied to PyTorch model: {} -" . format ( ', ' . join ( tf_weights . keys ( ) ) ) ) logger . info ( "Remanining parameters/buffers from PyTorch model: {} -" . format ( ', ' . join ( remaining_params ) ) ) return model
short_docstring: Load tf checkpoints and standing statistics in a pytorch model
short_docstring_tokens: Load tf checkpoints and standing statistics in a pytorch model
comment: """ Load tf checkpoints and standing statistics in a pytorch model
      """$SEP$# Load weights from TF model$SEP$# Extract batch norm statistics from model if needed$SEP$# Build TF to PyTorch weights loading map$SEP$# logger.info("Loading TF weight {} with shape {}".format(name, array.shape))$SEP$# Load parameters$SEP$# Weights$SEP$# Convolutions$SEP$# Prepare SpectralNorm buffers by running one step of Spectral Norm (no need to train the model):$SEP$# Load batch norm statistics$SEP$#  Batchnorms