Load the weights from the tensorflow checkpoint
(inputbase, config)
| 67 | |
| 68 | |
| 69 | def load_tf_weights(inputbase, config): |
| 70 | """ |
| 71 | Load the weights from the tensorflow checkpoint |
| 72 | """ |
| 73 | weights_dict = dict() |
| 74 | |
| 75 | try: |
| 76 | reader = tf.train.NewCheckpointReader(inputbase) |
| 77 | tensor_dict = reader.get_variable_to_shape_map() |
| 78 | |
| 79 | # There might be training-related variables in the checkpoint that can be discarded |
| 80 | param_names = [key for key in sorted(tensor_dict) if "adam" not in key and "global_step" not in key and "pooler" not in key] |
| 81 | count = len(param_names) |
| 82 | TRT_LOGGER.log(TRT_LOGGER.INFO, "Found {:} entries in weight map".format(count)) |
| 83 | |
| 84 | for pn in param_names: |
| 85 | toks = pn.lower().split("/") |
| 86 | if "encoder" in pn: |
| 87 | assert ("layer" in pn) |
| 88 | l = (re.findall("\d+", pn))[0] |
| 89 | outname = "l{}_".format(l) + "_".join(toks[3:]) |
| 90 | else: |
| 91 | outname = "_".join(toks) |
| 92 | |
| 93 | tensor = reader.get_tensor(pn) |
| 94 | shape = tensor.shape |
| 95 | if pn.find("kernel") != -1: |
| 96 | weights_dict[outname + "_notrans"] = trt.Weights(np.ascontiguousarray(tensor).flatten()) |
| 97 | |
| 98 | TRT_LOGGER.log(TRT_LOGGER.VERBOSE, "Transposing {}\n".format(np)) |
| 99 | tensor = np.transpose(tensor) |
| 100 | |
| 101 | shape = tensor.shape |
| 102 | flat_tensor = tensor.flatten() |
| 103 | shape_str = "{} ".format(len(shape)) + " ".join([str(d) for d in shape]) |
| 104 | weights_dict[outname] = trt.Weights(flat_tensor) |
| 105 | |
| 106 | TRT_LOGGER.log(TRT_LOGGER.VERBOSE, "Original name: {:}, TensorRT name: {:}, shape: {:}".format(pn, outname, shape_str)) |
| 107 | |
| 108 | N = config.num_attention_heads |
| 109 | H = config.head_size |
| 110 | |
| 111 | additional_dict = dict() |
| 112 | for key, value in weights_dict.items(): |
| 113 | pos = key.find(BQ) |
| 114 | if pos != -1: |
| 115 | hidden_size = value.size |
| 116 | prefix = key[:pos] |
| 117 | |
| 118 | Bq_ = value |
| 119 | Bk_ = weights_dict[prefix + BK] |
| 120 | Bv_ = weights_dict[prefix + BV] |
| 121 | Wq_ = weights_dict[prefix + WQ] |
| 122 | Wk_ = weights_dict[prefix + WK] |
| 123 | Wv_ = weights_dict[prefix + WV] |
| 124 | |
| 125 | mat_size = hidden_size * hidden_size |
| 126 | wcount = 3 * mat_size |