Run the bert transformer layer by FasterTransformer. Args: inputs: A tf.Tensor with shape [batch_size, seq_len, hidden_dimension]. The inputs tensor of encoder. The rank must be 3. encoder_args: The arguments for encoder. The details are in the class "Transfo
(inputs,
encoder_args,
encoder_vars_dict,
sequence_length)
| 201 | |
| 202 | |
| 203 | def ft_encoder_opennmt(inputs, |
| 204 | encoder_args, |
| 205 | encoder_vars_dict, |
| 206 | sequence_length): |
| 207 | |
| 208 | ''' |
| 209 | Run the bert transformer layer by FasterTransformer. |
| 210 | Args: |
| 211 | inputs: A tf.Tensor with shape [batch_size, seq_len, hidden_dimension]. |
| 212 | The inputs tensor of encoder. The rank must be 3. |
| 213 | encoder_args: The arguments for encoder. The details are in the class "TransformerArgument" of common.py |
| 214 | attention_mask: A tf.Tensor. The attention mask for self attention. |
| 215 | encoder_vars_dict: A dict of tf.Tensor or numpy array. |
| 216 | The variables for encoder. They can be either some tensor or some numpy array. |
| 217 | The key is the name of the tensor, like 'layer_0/attention/self/query/kernel:0'. |
| 218 | The value is the corresponding tensor or numpy array |
| 219 | sequence_length: A tf.Tensor or numpy array with shape [batch_size]. |
| 220 | The sequence length of the sentences |
| 221 | Outputs: |
| 222 | outputs: A tensor with shape [batch_size, seq_len, hidden_dimension]. |
| 223 | The results of encoder. |
| 224 | ''' |
| 225 | |
| 226 | q_w_list = [] |
| 227 | q_b_list = [] |
| 228 | k_w_list = [] |
| 229 | k_b_list = [] |
| 230 | v_w_list = [] |
| 231 | v_b_list = [] |
| 232 | |
| 233 | for i in range(encoder_args.num_layer): |
| 234 | q_w, k_w, v_w = tf.split(encoder_vars_dict['transformer/encoder/layer_%d/multi_head/conv1d/kernel:0' % i], 3, axis=-1) |
| 235 | q_w_list.append(q_w) |
| 236 | k_w_list.append(k_w) |
| 237 | v_w_list.append(v_w) |
| 238 | |
| 239 | q_b, k_b, v_b = tf.split(encoder_vars_dict['transformer/encoder/layer_%d/multi_head/conv1d/bias:0' % i], 3, axis=-1) |
| 240 | q_b_list.append(q_b) |
| 241 | k_b_list.append(k_b) |
| 242 | v_b_list.append(v_b) |
| 243 | |
| 244 | input_shape = get_shape_list(inputs, expected_rank=3) |
| 245 | seq_length = input_shape[1] |
| 246 | inputs *= encoder_args.hidden_dim**0.5 |
| 247 | position_encoder = SinusoidalPositionEncoder() |
| 248 | inputs = position_encoder(inputs, position=tf.range(seq_length)) |
| 249 | |
| 250 | transformer_op_module = tf.load_op_library(os.path.join('./lib/libtf_encoder.so')) |
| 251 | tf_datatype = inputs.dtype |
| 252 | outputs = transformer_op_module.encoder( |
| 253 | inputs, |
| 254 | inputs, |
| 255 | sequence_length, |
| 256 | [tf.cast(encoder_vars_dict['transformer/encoder/layer_%d/multi_head/LayerNorm/beta:0' % id], tf_datatype) for id in range(encoder_args.num_layer)], |
| 257 | [tf.cast(encoder_vars_dict['transformer/encoder/layer_%d/multi_head/LayerNorm/gamma:0' % id], tf_datatype) for id in range(encoder_args.num_layer)], |
| 258 | q_w_list, q_b_list, |
| 259 | k_w_list, k_b_list, |
| 260 | v_w_list, v_b_list, |
no test coverage detected