# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # refrenece : https://github.com/PaddlePaddle/PaddleNLP/tree/develop/examples/language_model/ernie import os import copy import argparse from contextlib import contextmanager from functools import partial import numpy as np import paddle import paddle.static import paddle.fluid as fluid import paddle.fluid.layers as layers import paddle.fluid.compiler as compiler paddle.enable_static() SEED = 2021 INT_DTYPE = None # ernie related block ernie_config = { "emb_size": 128, "emb_mapping_in": False, "hidden_size": 192, "num_hidden_layers": 2, "n_layer_per_block": 2, "num_attention_heads": 12, "vocab_size": 300, "max_position_embeddings": 512, "sent_type_vocab_size": 4, "task_type_vocab_size": 16, "hidden_act": "gelu", "hidden_dropout_prob": 0.0, "attention_probs_dropout_prob": 0.0, "preln": False, "pre_encoder_cmd": "n", "preprocess_cmd": "", "postprocess_cmd": "an", "epsilon": 1e-12, "initializer_range": 0.02, "seq_len": 32 } def gelu(x): """Gaussian Error Linear Unit. This is a smoother version of the RELU. Original paper: https://arxiv.org/abs/1606.08415 Args: x: float Tensor to perform activation. Returns: `x` with the GELU activation applied. """ cdf = 0.5 * (1.0 + fluid.layers.tanh( (np.sqrt(2.0 / np.pi) * (x + 0.044715 * fluid.layers.pow(x, 3.0))))) return x * cdf def pre_post_process_layer(prev_out, out, process_cmd, dropout_rate=0., epsilon=1e-12, name=''): """ Add residual connection, layer normalization and droput to the out tensor optionally according to the value of process_cmd. This will be used before or after multi-head attention and position-wise feed-forward networks. """ for cmd in process_cmd: if cmd == "a": # add residual connection out = out + prev_out if prev_out else out elif cmd == "n": # add layer normalization out = layers.layer_norm( out, begin_norm_axis=len(out.shape) - 1, param_attr=fluid.ParamAttr( name=name + '_layer_norm_scale', initializer=fluid.initializer.Constant(1.)), bias_attr=fluid.ParamAttr( name=name + '_layer_norm_bias', initializer=fluid.initializer.Constant(0.)), epsilon=epsilon) elif cmd == "d": # add dropout if dropout_rate: out = layers.dropout( out, dropout_prob=dropout_rate, dropout_implementation="upscale_in_train", is_test=False) return out pre_process_layer = partial(pre_post_process_layer, None) post_process_layer = pre_post_process_layer def positionwise_feed_forward(x, d_inner_hid, d_hid, dropout_rate, hidden_act, param_initializer=None, name='ffn'): """ Position-wise Feed-Forward Networks. This module consists of two linear transformations with a ReLU activation in between, which is applied to each position separately and identically. """ #assert hidden_act == 'gelu.approximate' hidden = layers.fc(input=x, size=d_inner_hid, num_flatten_dims=2, act=None, param_attr=fluid.ParamAttr( name=name + '_fc_0.w_0', initializer=param_initializer), bias_attr=name + '_fc_0.b_0') hidden = gelu(hidden) if dropout_rate: hidden = layers.dropout( hidden, dropout_prob=dropout_rate, dropout_implementation="upscale_in_train", is_test=False) out = layers.fc(input=hidden, size=d_hid, num_flatten_dims=2, param_attr=fluid.ParamAttr( name=name + '_fc_1.w_0', initializer=param_initializer), bias_attr=name + '_fc_1.b_0') return out def multi_head_attention(queries, keys, values, attn_bias, d_key, d_value, d_model, n_head=1, dropout_rate=0., cache=None, param_initializer=None, name='multi_head_att'): """ Multi-Head Attention. Note that attn_bias is added to the logit before computing softmax activiation to mask certain selected positions so that they will not considered in attention weights. """ keys = queries if keys is None else keys values = keys if values is None else values def __compute_qkv(queries, keys, values, n_head, d_key, d_value): """ Add linear projection to queries, keys, and values. """ q = layers.fc(input=queries, size=d_key * n_head, num_flatten_dims=2, param_attr=fluid.ParamAttr( name=name + '_query_fc.w_0', initializer=param_initializer), bias_attr=name + '_query_fc.b_0') k = layers.fc(input=keys, size=d_key * n_head, num_flatten_dims=2, param_attr=fluid.ParamAttr( name=name + '_key_fc.w_0', initializer=param_initializer), bias_attr=name + '_key_fc.b_0') v = layers.fc(input=values, size=d_value * n_head, num_flatten_dims=2, param_attr=fluid.ParamAttr( name=name + '_value_fc.w_0', initializer=param_initializer), bias_attr=name + '_value_fc.b_0') return q, k, v def __split_heads(x, n_head): """ Reshape the last dimension of inpunt tensor x so that it becomes two dimensions and then transpose. Specifically, input a tensor with shape [bs, max_sequence_length, n_head * hidden_dim] then output a tensor with shape [bs, n_head, max_sequence_length, hidden_dim]. """ hidden_size = x.shape[-1] # The value 0 in shape attr means copying the corresponding dimension # size of the input as the output dimension size. reshaped = layers.reshape( x=x, shape=[0, 0, n_head, hidden_size // n_head], inplace=False) # permuate the dimensions into: # [batch_size, n_head, max_sequence_len, hidden_size_per_head] return layers.transpose(x=reshaped, perm=[0, 2, 1, 3]) def __combine_heads(x): """ Transpose and then reshape the last two dimensions of inpunt tensor x so that it becomes one dimension, which is reverse to __split_heads. """ if len(x.shape) == 3: return x if len(x.shape) != 4: raise ValueError("Input(x) should be a 4-D Tensor.") trans_x = layers.transpose(x, perm=[0, 2, 1, 3]) # The value 0 in shape attr means copying the corresponding dimension # size of the input as the output dimension size. return layers.reshape( x=trans_x, shape=[0, 0, trans_x.shape[2] * trans_x.shape[3]], inplace=False) def scaled_dot_product_attention(q, k, v, attn_bias, d_key, dropout_rate): """ Scaled Dot-Product Attention """ scaled_q = layers.scale(x=q, scale=d_key**-0.5) product = layers.matmul(x=scaled_q, y=k, transpose_y=True) if attn_bias: product += attn_bias weights = layers.softmax(product) if dropout_rate: weights = layers.dropout( weights, dropout_prob=dropout_rate, dropout_implementation="upscale_in_train", is_test=False) out = layers.matmul(weights, v) return out q, k, v = __compute_qkv(queries, keys, values, n_head, d_key, d_value) if cache is not None: # use cache and concat time steps # Since the inplace reshape in __split_heads changes the shape of k and # v, which is the cache input for next time step, reshape the cache # input from the previous time step first. k = cache["k"] = layers.concat( [layers.reshape( cache["k"], shape=[0, 0, d_model]), k], axis=1) v = cache["v"] = layers.concat( [layers.reshape( cache["v"], shape=[0, 0, d_model]), v], axis=1) q = __split_heads(q, n_head) k = __split_heads(k, n_head) v = __split_heads(v, n_head) ctx_multiheads = scaled_dot_product_attention(q, k, v, attn_bias, d_key, dropout_rate) out = __combine_heads(ctx_multiheads) # Project back to the model size. proj_out = layers.fc(input=out, size=d_model, num_flatten_dims=2, param_attr=fluid.ParamAttr( name=name + '_output_fc.w_0', initializer=param_initializer), bias_attr=name + '_output_fc.b_0') return proj_out def encoder_layer(enc_input, attn_bias, n_head, d_key, d_value, d_model, d_inner_hid, prepostprocess_dropout, attention_dropout, relu_dropout, hidden_act, preprocess_cmd="n", postprocess_cmd="da", param_initializer=None, name='', epsilon=1e-12): """The encoder layers that can be stacked to form a deep encoder. This module consits of a multi-head (self) attention followed by position-wise feed-forward networks and both the two components companied with the post_process_layer to add residual connection, layer normalization and droput. """ attn_output = multi_head_attention( enc_input, None, None, attn_bias, d_key, d_value, d_model, n_head, attention_dropout, param_initializer=param_initializer, name=name + '_multi_head_att') attn_output = post_process_layer( enc_input, attn_output, 'an', prepostprocess_dropout, name=name + '_post_att', epsilon=epsilon) ffd_output = positionwise_feed_forward( attn_output, d_inner_hid, d_model, relu_dropout, hidden_act, param_initializer=param_initializer, name=name + '_ffn') post_output = post_process_layer( attn_output, ffd_output, 'an', prepostprocess_dropout, name=name + '_post_ffn', epsilon=epsilon) return post_output def encoder_inner_share(enc_input, attn_bias, n_head, d_key, d_value, d_model, d_inner_hid, prepostprocess_dropout, attention_dropout, relu_dropout, hidden_act, preprocess_cmd, postprocess_cmd, epsilon, param_initializer=None, name='', n_layer_per_block=1): """ The encoder_inner_share is composed of n_layer_per_block layers returned by calling encoder_layer. """ for i in range(n_layer_per_block): enc_output = encoder_layer( enc_input, attn_bias, n_head, d_key, d_value, d_model, d_inner_hid, prepostprocess_dropout, attention_dropout, relu_dropout, hidden_act, preprocess_cmd, postprocess_cmd, param_initializer=param_initializer, name=name + '_layer_' + str(i), epsilon=epsilon) enc_input = enc_output return enc_output def encoder(enc_input, attn_bias, n_layer, n_head, d_key, d_value, d_model, d_inner_hid, prepostprocess_dropout, attention_dropout, relu_dropout, hidden_act, preprocess_cmd, postprocess_cmd, epsilon, n_layer_per_block, param_initializer=None, name='', preln=False): """ The encoder is composed of a stack of identical layers returned by calling encoder_layer . """ for _ in range(n_layer // n_layer_per_block): attn_bias.stop_gradient = True attn_bias.persistable = False enc_output = encoder_inner_share( enc_input, attn_bias, n_head, d_key, d_value, d_model, d_inner_hid, prepostprocess_dropout, attention_dropout, relu_dropout, hidden_act, preprocess_cmd, postprocess_cmd, epsilon, param_initializer=param_initializer, name=name, n_layer_per_block=n_layer_per_block) enc_input = enc_output if preln: enc_output = post_process_layer( None, enc_output, 'n', prepostprocess_dropout, name='post_encoder', epsilon=epsilon) enc_output = pre_process_layer( enc_output, preprocess_cmd, prepostprocess_dropout, name="post_encoder", epsilon=epsilon) return enc_output class ErnieModel(object): def __init__(self, src_ids, sent_ids, pos_ids, input_mask, config): self._emb_size = config['emb_size'] if config[ 'emb_mapping_in'] else config['hidden_size'] self._hidden_size = config['hidden_size'] self._n_layer = config['num_hidden_layers'] self._n_head = config['num_attention_heads'] self._voc_size = config['vocab_size'] self._max_position_seq_len = config['max_position_embeddings'] self._sent_types = config['sent_type_vocab_size'] self._task_types = config['task_type_vocab_size'] self._hidden_act = config['hidden_act'] self._prepostprocess_dropout = config['hidden_dropout_prob'] self._attention_dropout = config['attention_probs_dropout_prob'] self.config = config self.preln = config['preln'] if 'preln' in config.keys() else False self.pre_encoder_cmd = "" if self.preln else self.config[ 'pre_encoder_cmd'] self._word_emb_name = "word_embedding" self._pos_emb_name = "pos_embedding" self._sent_emb_name = "sent_embedding" self._task_emb_name = "task_embedding" self._dtype = "float32" self._emb_dtype = "float32" # Initialize all weigths by truncated normal initializer, and all biases # will be initialized by constant zero by default. self._param_initializer = fluid.initializer.TruncatedNormal( scale=config['initializer_range']) self.src_ids = src_ids self.sent_ids = sent_ids self.pos_ids = pos_ids self.input_mask = input_mask ''' _build_position_ids: range op doesn't support _build_input_mask: logic_not op doesn't support ''' self._build_model() def _build_model(self, emb=None): with fluid.ipu_shard(ipu_index=0, ipu_stage=0): # padding id in vocabulary must be set to 0 self.emb_out = fluid.layers.embedding( input=self.src_ids, size=[self._voc_size, self._emb_size], dtype=self._emb_dtype, param_attr=fluid.ParamAttr( name=self._word_emb_name, initializer=self._param_initializer), is_sparse=False) self.position_emb_out = fluid.layers.embedding( input=self.pos_ids, size=[self._max_position_seq_len, self._emb_size], dtype=self._emb_dtype, param_attr=fluid.ParamAttr( name=self._pos_emb_name, initializer=self._param_initializer)) self.sent_emb_out = fluid.layers.embedding( self.sent_ids, size=[self._sent_types, self._emb_size], dtype=self._emb_dtype, param_attr=fluid.ParamAttr( name=self._sent_emb_name, initializer=self._param_initializer)) sum_emb = self.emb_out + self.position_emb_out + self.sent_emb_out sum_emb = pre_process_layer( sum_emb, self.config['pre_encoder_cmd'], self._prepostprocess_dropout, name='pre_encoder', epsilon=self.config['epsilon']) if self.config['emb_mapping_in']: sum_emb = fluid.layers.fc( input=sum_emb, num_flatten_dims=2, size=self._hidden_size, param_attr=fluid.ParamAttr( name='emb_hidden_mapping', initializer=self._param_initializer), bias_attr='emb_hidden_mapping_bias') self_attn_mask = fluid.layers.matmul( x=self.input_mask, y=self.input_mask, transpose_y=True) self_attn_mask = fluid.layers.scale( x=self_attn_mask, scale=10000.0, bias=-1.0, bias_after_scale=False) with fluid.ipu_shard(ipu_index=1, ipu_stage=1): n_head_self_attn_mask = fluid.layers.stack( x=[self_attn_mask] * self._n_head, axis=1) # [bs, _n_head, seqlen, seq_len] n_head_self_attn_mask.stop_gradient = True self._enc_out = encoder( enc_input=sum_emb, attn_bias=n_head_self_attn_mask, n_layer=self._n_layer, n_head=self._n_head, d_key=self._hidden_size // self._n_head, d_value=self._hidden_size // self._n_head, d_model=self._hidden_size, d_inner_hid=self._hidden_size * 4, prepostprocess_dropout=self._prepostprocess_dropout, attention_dropout=self._attention_dropout, relu_dropout=0, hidden_act=self._hidden_act, preprocess_cmd=self.config['preprocess_cmd'], postprocess_cmd=self.config['postprocess_cmd'], param_initializer=self._param_initializer, name='encoder', epsilon=self.config['epsilon'], n_layer_per_block=self.config['n_layer_per_block'], preln=self.preln) def _build_position_ids(self): d_shape = fluid.layers.shape(self.src_ids) d_seqlen = d_shape[1] d_batch = d_shape[0] position_ids = fluid.layers.reshape( fluid.layers.range( 0, d_seqlen, 1, dtype='int32'), [1, d_seqlen, 1], inplace=False) position_ids = fluid.layers.expand(position_ids, [d_batch, 1, 1]) position_ids = fluid.layers.cast(position_ids, INT_DTYPE) position_ids.stop_gradient = True return position_ids def _build_input_mask(self): zero = fluid.layers.fill_constant([1], dtype=INT_DTYPE, value=0) input_mask = fluid.layers.logical_not( fluid.layers.equal(self.src_ids, zero)) # assume pad id == 0 input_mask = fluid.layers.cast(input_mask, 'float32') input_mask.stop_gradient = True return input_mask def get_sequence_output(self): return self._enc_out def get_pooled_output(self): """Get the first feature of each sequence for classification""" next_sent_feat = fluid.layers.slice( input=self._enc_out, axes=[1], starts=[0], ends=[1]) next_sent_feat = fluid.layers.fc( input=next_sent_feat, size=self._hidden_size, act="tanh", param_attr=fluid.ParamAttr( name="pooled_fc.w_0", initializer=self._param_initializer), bias_attr="pooled_fc.b_0") return next_sent_feat def get_next_sentence_output(self, labels): next_sent_feat = self.get_pooled_output() next_sent_fc_out = fluid.layers.fc( input=next_sent_feat, num_flatten_dims=1, size=33, param_attr=fluid.ParamAttr( name="next_sent_fc.w_0", initializer=self._param_initializer), bias_attr="next_sent_fc.b_0") next_sent_fc_out = fluid.layers.reshape( next_sent_fc_out, [-1, 33], inplace=False) #next_sent_loss, next_sent_softmax = fluid.layers.softmax_with_cross_entropy( # logits=next_sent_fc_out, label=labels, return_softmax=True) next_sent_softmax = fluid.layers.softmax(next_sent_fc_out) next_sent_loss = fluid.layers.cross_entropy(next_sent_softmax, labels) next_sent_acc = fluid.layers.accuracy( input=next_sent_softmax, label=labels) mean_next_sent_loss = fluid.layers.mean(next_sent_loss, "mean_next_sent_loss") return next_sent_acc, mean_next_sent_loss def get_lm_output(self, mask_label, mask_pos): """Get the loss & accuracy for pretraining""" mask_pos = fluid.layers.cast(x=mask_pos, dtype='int32') # extract the first token feature in each sentence reshaped_emb_out = fluid.layers.reshape( x=self._enc_out, shape=[-1, self._hidden_size]) # extract masked tokens' feature mask_feat = fluid.layers.gather(input=reshaped_emb_out, index=mask_pos) if self._dtype == "float16": mask_feat = fluid.layers.cast(x=mask_feat, dtype=self._emb_dtype) # transform: fc if self._hidden_act == 'gelu' or self._hidden_act == 'gelu.precise': _hidden_act = 'gelu' else: _hidden_act = None mask_trans_feat = fluid.layers.fc( input=mask_feat, size=self._emb_size, act=_hidden_act, param_attr=fluid.ParamAttr( name='mask_lm_trans_fc.w_0', initializer=self._param_initializer), bias_attr=fluid.ParamAttr(name='mask_lm_trans_fc.b_0')) if self._hidden_act == 'gelu' or self._hidden_act == 'gelu.precise': pass else: mask_trans_feat = gelu(mask_trans_feat) # transform: layer norm mask_trans_feat = fluid.layers.layer_norm( mask_trans_feat, begin_norm_axis=len(mask_trans_feat.shape) - 1, param_attr=fluid.ParamAttr( name='mask_lm_trans_layer_norm_scale', initializer=fluid.initializer.Constant(1.)), bias_attr=fluid.ParamAttr( name='mask_lm_trans_layer_norm_bias', initializer=fluid.initializer.Constant(0.)), epsilon=self.config['epsilon']) mask_lm_out_bias_attr = fluid.ParamAttr( name="mask_lm_out_fc.b_0", initializer=fluid.initializer.Constant(value=0.0)) fc_out = fluid.layers.fc(input=mask_trans_feat, size=self._voc_size, param_attr=fluid.ParamAttr( name="mask_lm_out_fc.w_0", initializer=self._param_initializer), bias_attr=mask_lm_out_bias_attr) #mask_lm_loss = fluid.layers.softmax_with_cross_entropy( # logits=fc_out, label=mask_label) mask_lm_softmax = fluid.layers.softmax(fc_out) mask_lm_loss = fluid.layers.cross_entropy(mask_lm_softmax, mask_label) mean_mask_lm_loss = fluid.layers.mean( mask_lm_loss, name="mean_mask_lm_loss") return mask_lm_loss, mean_mask_lm_loss def get_task_output(self, task, task_labels): task_fc_out = fluid.layers.fc(input=self.next_sent_feat, size=task["num_labels"], param_attr=fluid.ParamAttr( name=task["task_name"] + "_fc.w_0", initializer=self._param_initializer), bias_attr=task["task_name"] + "_fc.b_0") #task_loss, task_softmax = fluid.layers.softmax_with_cross_entropy( # logits=task_fc_out, label=task_labels, return_softmax=True) task_softmax = fluid.layers.softmax(task_fc_out) task_loss = fluid.layers.cross_entropy(task_softmax, task_labels) task_acc = fluid.layers.accuracy(input=task_softmax, label=task_labels) mean_task_loss = fluid.layers.mean(task_loss) return mean_task_loss, task_acc if __name__ == "__main__": parser = argparse.ArgumentParser(__doc__) parser.add_argument( "--run_on_ipu", type=bool, default=True, help="Run model with IPU") parser.add_argument( "--is_training", type=bool, default=True, help="Train of inference") parser.add_argument( "--num_ipus", type=int, default=2, help="Number of ipus") parser.add_argument( "--enable_pipelining", type=bool, default=False, help="Pipelining") parser.add_argument( "--save_model", type=bool, default=False, help="Save model or not") parser.add_argument( "--model_path", type=str, default="ernie", help="Save model to where") parser.add_argument( "--model_name", type=str, default="ernie", help="Save model name") parser.add_argument( "--ipu_run_steps", type=int, default=10, help="Number steps exe.run()") parser.add_argument( "--export_ops", type=bool, default=False, help="Export ops to ops.txt") parser.add_argument( "--export_ipu_idx", type=bool, default=False, help="Export op-idx pair") args = parser.parse_args() # set random seed np.random.seed(SEED) paddle.static.default_startup_program().random_seed = SEED paddle.static.default_main_program().random_seed = SEED # IPU doesn't support int64, so we change here INT_DTYPE = "int32" if args.run_on_ipu else "int64" # paddle input placeholder, batch_size = 1 micro_bs = 1 seq_len = ernie_config["seq_len"] input_shape = [micro_bs, seq_len, 1] input_fields = { 'names': [ 'src_ids', 'sent_ids', 'pos_ids', 'input_mask', 'mask_label', 'mask_pos' ], 'shapes': [ input_shape, input_shape, input_shape, input_shape, [micro_bs, 1], [micro_bs, 1] ], 'dtypes': [INT_DTYPE, INT_DTYPE, INT_DTYPE, 'float32', INT_DTYPE, INT_DTYPE], 'range': [[0, seq_len], [0, 4], [0, seq_len], None, [0, seq_len], [0, seq_len]], 'lod_levels': [0, 0, 0, 0, 0, 0], } inputs = [ fluid.data( name=input_fields['names'][i], shape=input_fields['shapes'][i], dtype=input_fields['dtypes'][i], lod_level=input_fields['lod_levels'][i]) for i in range(len(input_fields['names'])) ] # total_samples: assum disable pipelining batches_per_step = 1 if args.enable_pipelining: batches_per_step = \ ((args.num_ipus+1) if args.is_training else args.num_ipus) total_samples = args.ipu_run_steps * batches_per_step total_steps = args.ipu_run_steps if not args.run_on_ipu: # run on cpu total_steps = total_samples // micro_bs # synthetic data np_inputs = [] for i in range(len(input_fields['names'])): field_name = input_fields['names'][i] if field_name == 'input_mask': src_ids = np_inputs[0] dtype = input_fields['dtypes'][i] data = np.where(src_ids > 0, np.ones_like(src_ids), np.zeros_like(src_ids)).astype(dtype) else: shape = copy.copy(input_fields['shapes'][i]) shape[0] = total_samples min_val, max_val = input_fields['range'][i] data = np.random.randint( min_val, max_val, shape, dtype=input_fields['dtypes'][i]) np_inputs.append(data) # paddle input placeholder (src_ids, sent_ids, pos_ids, input_mask, mask_label, mask_pos) = inputs # ernie model ernie = ErnieModel(src_ids, sent_ids, pos_ids, input_mask, ernie_config) fetch_node = ernie.get_sequence_output() if args.is_training: with fluid.ipu_shard(ipu_index=1, ipu_stage=1): _, mean_mask_lm_loss = ernie.get_lm_output(mask_label, mask_pos) fetch_node = mean_mask_lm_loss adam = paddle.optimizer.Adam(learning_rate=1e-2) adam.minimize(mean_mask_lm_loss) # place = paddle.CPUPlace() if args.run_on_ipu: place = paddle.IPUPlace() else: place = paddle.CPUPlace() executor = paddle.static.Executor(place) # feed & fetch list if args.is_training: feed_list = input_fields['names'] else: feed_list = input_fields['names'][:4] fetch_list = [fetch_node.name] # program startup_prog = paddle.static.default_startup_program() executor.run(startup_prog) main_prog = paddle.static.default_main_program() paddle.static.save(main_prog, "model/ernie") paddle.static.load(main_prog, "model/ernie") if args.run_on_ipu: ipu_strategy = compiler.get_ipu_strategy() ipu_strategy.num_ipus = args.num_ipus ipu_strategy.enable_manual_shard = args.num_ipus > 1 ipu_strategy.enable_pipelining = args.enable_pipelining if args.enable_pipelining: if args.is_training: ipu_strategy.batches_per_step = args.num_ipus + 1 else: ipu_strategy.batches_per_step = args.num_ipus ipu_strategy.is_training = args.is_training ipu_compiler = compiler.IPUCompiledProgram( main_prog, ipu_strategy=ipu_strategy) program = ipu_compiler.compile(feed_list, fetch_list) else: program = main_prog # executor run results = [] for i in range(total_steps): start = i * (batches_per_step if args.run_on_ipu else 1) end = start + (batches_per_step if args.run_on_ipu else 1) feed_dict = { src_ids.name: np_inputs[0][start:end], sent_ids.name: np_inputs[1][start:end], pos_ids.name: np_inputs[2][start:end], input_mask.name: np_inputs[3][start:end] } if args.is_training: feed_dict[mask_label.name] = np_inputs[4][start:end] feed_dict[mask_pos.name] = np_inputs[5][start:end] res = executor.run(program, feed=feed_dict, fetch_list=[fetch_node]) results.append(res) paddle.static.save(main_prog, "model/ernie") results = np.asarray(results).flatten() if results.size > 32: results = results[-32:] print(results) if args.save_model: full_name = args.model_path + '/' + args.model_name if args.is_training: fluid.save(program=main_prog, model_path=full_name) else: with fluid.ipu_shard(ipu_index=1, ipu_stage=1): paddle.static.save_inference_model( full_name, [src_ids, sent_ids, pos_ids, input_mask], [fetch_node], executor) if args.export_ops: op_type_list = [] for op in main_prog.global_block().ops: op_type_list.append(op.desc.type()) with open("ops.txt", "w") as fp: for op_type in set(op_type_list): fp.write(op_type + os.linesep) if args.export_ipu_idx: op_ipu_idx_list = [] for op in main_prog.global_block().ops: if op._is_backward_op(): continue op_ipu_idx_pair = [op.desc.type()] if op.desc.has_attr("ipu_index"): op_ipu_idx_pair.append(op.desc.attr("ipu_index")) else: op_ipu_idx_pair.append(-1) # not assign ipu_index op_ipu_idx_list.append(op_ipu_idx_pair) op_ipu_idx_list.sort(key=lambda item: item[-1]) with open("ops_ipu_idx.txt", "w") as fp: for op_ipu_idx_pair in op_ipu_idx_list: fp.write(str(op_ipu_idx_pair) + os.linesep)