code: def _get_assignment_map_from_checkpoint(
      sess, all_variables, init_checkpoint
  ):
      """
      Load pretrained parameters to texar model
      """
  
      assignment_map = {}
  
      reader = tf.train.NewCheckpointReader(init_checkpoint)
      var_names_list = reader.get_variable_to_shape_map().keys()
      ckpt_names_vs_vals = {}
      for var_name in var_names_list:
          ckpt_names_vs_vals[var_name] = reader.get_tensor(var_name)
  
      def _assign_by_name(sess, tensor_name, data):
          local_tensor = [
              var for var in all_variables if tensor_name in var.name
          ][0]
          sess.run(tf.assign(local_tensor, data))
  
      def _get_tensor_by_name(tensor_name):
          local_tensor = [
              var for var in all_variables if tensor_name in var.name
          ][0]
          return local_tensor
  
      for idx, ckpt_tensor_name in enumerate(ckpt_names_vs_vals):
          processing = (idx + 1.0) / len(ckpt_names_vs_vals.keys())
          sys.stdout.write(
              "\rLoading checkpoint: {:.1%}".format(processing)
          )
          sys.stdout.flush()
  
          ckpt_tensor_name_feature = ""
          if len(ckpt_tensor_name.split("/")) > 2:
              ckpt_tensor_name_feature = "/".join(
                  ckpt_tensor_name.split("/")[2:]
              )
          if ckpt_tensor_name_feature == "attn/c_attn/w":
              layer_num = int(ckpt_tensor_name.split("/")[1][1:])
              template = (
                  "transformer_decoder/layer_{}/self_attention/"
                  "multihead_attention/{}/kernel"
              )
              local_tensor_name_q_w = template.format(
                  layer_num, "query"
              )
              local_tensor_name_k_w = template.format(layer_num, "key")
              local_tensor_name_v_w = template.format(
                  layer_num, "value"
              )
  
              data = ckpt_names_vs_vals[ckpt_tensor_name]
              assert data.shape[2] % 3 == 0, (
                  "tensor 'attn/c_attn/w' " "shape is not dividable"
              )
              index_w = data.shape[2] // 3
              q_w = data[:, :, :index_w]
              k_w = data[:, :, index_w : 2 * index_w]
              v_w = data[:, :, 2 * index_w :]
              _assign_by_name(
                  sess, local_tensor_name_q_w, np.squeeze(q_w)
              )
              _assign_by_name(
                  sess, local_tensor_name_k_w, np.squeeze(k_w)
              )
              _assign_by_name(
                  sess, local_tensor_name_v_w, np.squeeze(v_w)
              )
  
          elif ckpt_tensor_name_feature == "attn/c_attn/b":
              layer_num = int(ckpt_tensor_name.split("/")[1][1:])
              template = (
                  "transformer_decoder/layer_{}/self_attention/"
                  "multihead_attention/{}/bias"
              )
              local_tensor_name_q_b = template.format(
                  layer_num, "query"
              )
              local_tensor_name_k_b = template.format(layer_num, "key")
              local_tensor_name_v_b = template.format(
                  layer_num, "value"
              )
  
              data = ckpt_names_vs_vals[ckpt_tensor_name]
              assert data.shape[0] % 3 == 0, (
                  "tensor 'attn/c_attn/b'" " shape is not dividable"
              )
              index_b = data.shape[0] // 3
              q_b = data[:index_b]
              k_b = data[index_b : 2 * index_b]
              v_b = data[2 * index_b :]
              _assign_by_name(sess, local_tensor_name_q_b, q_b)
              _assign_by_name(sess, local_tensor_name_k_b, k_b)
              _assign_by_name(sess, local_tensor_name_v_b, v_b)
  
          else:
              local_tensor_name = _map_tensor_names(ckpt_tensor_name)
              local_tensor = _get_tensor_by_name(local_tensor_name)
              assignment_map[ckpt_tensor_name] = local_tensor
  
      return assignment_map
  
file_path: texar/examples/gpt-2/utils/model_utils.py
imports: import sys
  import json
  import tensorflow
  import numpy
  from texar.tf import HParams
label_desc: Use image node_id 3 for calling a global function with example usage: _get_assignment_map_from_checkpoint(sess, all_variables, init_checkpoint) and returns: assignment_map, local_tensor