def fwd_fn(idx, label): np.random.seed(1) embedding_shape = (dataset_size, embedding_size) embedding_initializer = np.random.normal(0, 1, embedding_shape).astype( np.float32) weights_shape = (embedding_size, embedding_size) weights_initializer = np.random.normal(0, 1, weights_shape).astype(np.float32) with variable_scope.variable_scope("part1", use_resource=True): embedding = variable_scope.get_variable( "c", dtype=np.float32, initializer=embedding_initializer, trainable=True) weight = variable_scope.get_variable("w0", dtype=np.float32, initializer=weights_initializer, trainable=True) x = embedding_ops.embedding_lookup(embedding, idx) x = math_ops.matmul(x, weight) logits = math_ops.reduce_sum(x, axis=[-1]) loss = math_ops.reduce_mean( nn.sparse_softmax_cross_entropy_with_logits(logits=logits, labels=label)) return loss
def stage1(idx, label): with variable_scope.variable_scope("stage1", use_resource=True): embedding = variable_scope.get_variable( "c", shape=[10, 1216], dtype=np.float32, initializer=init_ops.constant_initializer(10.01), trainable=True) x = embedding_ops.embedding_lookup(embedding, idx) return x, label
def _build_decoder_embedding(self, input, seq_len, scope_name): ''' Args: input: 2D-tensor, (batch_size, maxlen_tgt), target word index seq_len: 1D-tensor, (batch_size,), valid target word length scope_name: scope name Returns: x: 3D-tensor, (batch_size, maxlen_tgt, attention_dim) loss_mask: 3D-tensor, (batch_size, maxlen_tgt-1, 1) att_mask: 4D-tensor, (1, 1, maxlen_tgt-1, maxlen_tgt-1), constant ''' with tf.compat.v1.variable_scope(scope_name, dtype=self.dtype, use_resource=True) as scope: loss_mask = tf.sequence_mask(seq_len, maxlen=self.config['maxlen_tgt'] - 1, dtype=scope.dtype) loss_mask = tf.expand_dims(loss_mask, axis=2) embedding = tf.compat.v1.get_variable( "embedding_table", [self.config['vocab_size'], self.config['adim']], scope.dtype, initializer=tf.initializers.random_uniform(minval=0, maxval=1.0, dtype=scope.dtype), trainable=self.training, regularizer=self.kernel_regularizer) # x = tf.nn.embedding_lookup(embedding, input) x = embedding_ops.embedding_lookup(embedding, input) # position embedding _, max_len, dmodel = x.shape.as_list() pos_emb = self._build_pos_embedding(max_len, dmodel, reverse=False) x = math.sqrt(self.config['adim']) * x + pos_emb if self.training: if self.config['use_ipu_dropout']: x = ipu.rand_ops.dropout(x, rate=self.config['dropout_rate']) else: x = tf.nn.dropout(x, rate=self.config['dropout_rate']) # subsequent_mask index = tf.range(1, self.config['maxlen_tgt'], 1, dtype=tf.int32) index = tf.reshape(index, (1, 1, -1)) att_mask = tf.sequence_mask(index, dtype=scope.dtype) att_mask = (1.0 - att_mask) * self.mask_value return x, loss_mask, att_mask
def fwd_fn(idx, label): with variable_scope.variable_scope("part1", use_resource=True): embedding = variable_scope.get_variable( "c", shape=[10, 1216], dtype=np.float32, initializer=init_ops.constant_initializer(10.01), trainable=True) x = embedding_ops.embedding_lookup(embedding, idx) logits = math_ops.reduce_sum(x, axis=[-1]) loss = math_ops.reduce_mean( nn.sparse_softmax_cross_entropy_with_logits(logits=logits, labels=label)) return loss
def create_policy(*infeed_data): """Act according to current policy and generate action probability. """ dis_obs = list(infeed_data[:4]) cont_obs = list(infeed_data[4:8]) state_in = infeed_data[-1] # Look up embedding for all the discrete obs emb_lookup = [] with tf.variable_scope("popnn_lookup"): for index, obs in enumerate(dis_obs): emb_matrix = tf.get_variable( f'emb_matrix{index}', [DIS_OBS_CARDINALITY[index], DIS_OBS_EMB_SIZE[index]], DTYPE) emb_lookup.append( embedding_ops.embedding_lookup(emb_matrix, obs, name=f'emb_lookup{index}')) # Clip some continuous observations cont_obs[-1] = tf.clip_by_value(cont_obs[-1], -5.0, 5.0, name="clip") # Concat groups of observations obs_concat = [] for d_obs, c_obs in zip(emb_lookup, cont_obs): obs_concat.append(tf.concat([d_obs, c_obs], axis=3, name="concat_obs")) # Fully connected transformations num_output = 8 obs_concat[-1] = Dense(num_output, dtype=DTYPE)(obs_concat[-1]) # Reduce max obs_concat = [tf.reduce_max(obs, axis=2) for obs in obs_concat] # Final concat of all the observations lstm_input = tf.concat(obs_concat, axis=2, name="concat_all") # LSTM layer lstm_input = tf.transpose( lstm_input, perm=[1, 0, 2], name="pre_lstm_transpose") # PopnnLSTM uses time-major tensors lstm_cell = rnn_ops.PopnnLSTM(num_units=LSTM_HIDDEN_SIZE, dtype=DTYPE, partials_dtype=DTYPE, name="lstm") lstm_output, state_out = lstm_cell( lstm_input, training=True, initial_state=tf.nn.rnn_cell.LSTMStateTuple(state_in[:, 0], state_in[:, 1])) lstm_output = tf.transpose(lstm_output, perm=[1, 0, 2], name="post_lstm_transpose") logits = Dense(NUM_ACTIONS, name="logits", dtype=DTYPE)(lstm_output) log_prob = tf.nn.log_softmax(logits, name="prob") # make action selection op (outputs int actions, sampled from policy) actions = tf.random.categorical(logits=tf.reshape(logits, (-1, NUM_ACTIONS)), num_samples=1) actions = tf.reshape(actions, (args.batch_size, args.time_steps)) action_masks = tf.one_hot(actions, NUM_ACTIONS, dtype=DTYPE) action_prob = tf.reduce_sum(action_masks * log_prob, axis=-1) return action_prob