def __init__(self, strategy1, strategy2, strategy3): super().__init__() self.mul1 = P.Mul().set_strategy(strategy1) self.arg_min_with_value = P.ArgMinWithValue( keep_dims=True, axis=-1).set_strategy(strategy2) self.relu = P.ReLU().set_strategy(strategy3)
'block': P.Argmax(), 'desc_inputs': [[128, 32, 32, 64]], 'desc_bprop': [0], 'skip': ['backward']}), ('Argmin', { 'block': P.Argmin(), 'desc_inputs': [[128, 32, 32, 64]], 'desc_bprop': [1], 'skip': ['backward']}), ('ArgMaxWithValue', { 'block': P.ArgMaxWithValue(), 'desc_inputs': [[128, 32, 32, 64]], 'desc_bprop': [[1], [1]], 'skip': ['backward']}), ('ArgMinWithValue', { 'block': P.ArgMinWithValue(), 'desc_inputs': [[128, 32, 32, 64]], 'desc_bprop': [[1], [1]], 'skip': ['backward']}), ('Transpose_dim3', { 'block': P.Transpose(), 'desc_const': [(0, 2, 1)], 'desc_inputs': [[1, 2, 3]], 'desc_bprop': [[1, 3, 2]]}), ('Transpose_dim4', { 'block': P.Transpose(), 'desc_const': [(0, 1, 2, 3)], 'desc_inputs': [[1, 2, 3, 4]], 'desc_bprop': [[1, 2, 4, 3]]}), ('AddN', { 'block': NetForTupleInput(P.AddN()),
def __init__(self, batch_size, seq_length, vocab_size, decoder, beam_width=4, decoder_layers_nums=4, length_penalty_weight=0.6, cov_penalty_factor=0.1, hidden_size=1024, max_decode_length=64, sos_id=2, eos_id=3, compute_type=mstype.float32): super(BeamSearchDecoder, self).__init__() self.encoder_length = seq_length self.hidden_size = hidden_size self.batch_size = batch_size self.vocab_size = vocab_size self.beam_width = beam_width self.decoder_layers_nums = decoder_layers_nums self.length_penalty_weight = length_penalty_weight self.cov_penalty_factor = cov_penalty_factor self.max_decode_length = max_decode_length self.decoder = decoder self.add = P.TensorAdd() self.expand = P.ExpandDims() self.reshape = P.Reshape() self.shape_flat = (-1,) self.shape = P.Shape() self.zero_tensor = Tensor(np.zeros([batch_size, beam_width]), mstype.float32) self.ninf_tensor = Tensor(np.full([batch_size, beam_width], -INF), mstype.float32) self.select = P.Select() self.flat_shape = (batch_size, beam_width * vocab_size) self.topk = P.TopK(sorted=True) self.floor_div = P.FloorDiv() self.vocab_size_tensor = Tensor(self.vocab_size, mstype.int32) self.real_div = P.RealDiv() self.mod = Mod() self.equal = P.Equal() self.eos_ids = Tensor(np.full([batch_size, beam_width], eos_id), mstype.int32) beam_ids = np.tile(np.arange(beam_width).reshape((1, beam_width)), [batch_size, 1]) self.beam_ids = Tensor(beam_ids, mstype.int32) batch_ids = np.arange(batch_size * beam_width).reshape((batch_size, beam_width)) // beam_width self.batch_ids = Tensor(batch_ids, mstype.int32) self.concat = P.Concat(axis=-1) self.gather_nd = P.GatherNd() self.start = Tensor(0, dtype=mstype.int32) self.start_ids = Tensor(np.full([batch_size * beam_width, 1], sos_id), mstype.int32) self.init_seq = Tensor(np.full([batch_size, beam_width, self.max_decode_length], sos_id), mstype.int32) init_scores = np.tile(np.array([[0.] + [-INF] * (beam_width - 1)]), [batch_size, 1]) self.init_scores = Tensor(init_scores, mstype.float32) self.init_finished = Tensor(np.zeros([batch_size, beam_width], dtype=np.bool)) self.init_length = Tensor(np.zeros([batch_size, beam_width], dtype=np.int32)) self.length_penalty = LengthPenalty(weight=length_penalty_weight) self.one = Tensor(1, mstype.int32) self.prob_concat = P.Concat(axis=1) self.cast = P.Cast() self.decoder_hidden_state = Tensor(np.zeros([self.decoder_layers_nums, 2, self.batch_size * self.beam_width, hidden_size]), mstype.float32) self.zeros_scores = Tensor(np.zeros([batch_size, beam_width], dtype=np.float)) self.active_index = Tensor(np.ones([batch_size, beam_width], dtype=np.int32)) self.init_zeros = Tensor(np.zeros([batch_size, beam_width], dtype=np.int32)) self.init_ones = Tensor(np.ones([batch_size, beam_width], dtype=np.float32)) self.accu_attn_scores = Tensor(np.zeros([batch_size, beam_width, self.encoder_length], dtype=np.float32)) self.zeros = Tensor([0], mstype.int32) self.eos_tensor = Tensor(np.full([batch_size, beam_width, beam_width], eos_id), mstype.int32) self.ones_3d = Tensor(np.full([batch_size, beam_width, self.encoder_length], 1), mstype.float32) self.neg_inf_3d = Tensor(np.full([batch_size, beam_width, self.encoder_length], -INF), mstype.float32) self.zeros_3d = Tensor(np.full([batch_size, beam_width, self.encoder_length], 0), mstype.float32) self.zeros_2d = Tensor(np.full([batch_size * beam_width, self.encoder_length], 0), mstype.int32) self.argmin = P.ArgMinWithValue(axis=1) self.reducesum = P.ReduceSum() self.div = P.Div() self.shape_op = P.Shape() self.mul = P.Mul() self.log = P.Log() self.less = P.Less() self.tile = P.Tile() self.noteq = P.Neg() self.zeroslike = P.ZerosLike() self.greater_equal = P.GreaterEqual() self.sub = P.Sub()
def __init__(self, strategy1, strategy2, strategy3): super().__init__() self.mul1 = P.Mul().shard(strategy1) self.arg_min_with_value = P.ArgMinWithValue(keep_dims=False, axis=-1).shard(strategy2) self.mul2 = P.Mul().shard(strategy3)
def __init__(self, axis=0, keep_dims=False): super(NetArgminWithValue, self).__init__() self.argmin = P.ArgMinWithValue(axis=axis, keep_dims=keep_dims)