Tensorflow - Уменьшить издержки session.run ()

Question

В настоящее время я программирую небольшой учебный проект по подкреплению с tenorflow. Когда я профилировал тренировочный прогон, я заметил, что более 65% времени выполнения требуется только для выборки действий.

Примечание. В обучении с подкреплением вы должны выбирать действия на каждом шаге траектории. Вот почему не представляется возможным пробовать партиями.

Я запускаю этот график на графическом процессоре, поэтому я подозреваю, что издержки, связанные с вызовом метода session.run () и, следовательно, копированием данных из ОЗУ в графический процессор каждый раз, вызывают такие издержки.

Мой вопрос таков: есть ли способ уменьшить накладные расходы, связанные с вызовом метода session.run () столько раз? Можно ли запустить sample_op (прямой проход) на CPU и запустить обучение на GPU?

Заранее большое спасибо!

Это код моей политики RL:

class Policy:
def __init__(self, sess, state_size, action_size, lr, alpha_entropy,  epsilon):
    self.sess = sess
    self.action_size = action_size
    with tf.device("/cpu:0"):
        with tf.variable_scope("Policy"):
            self.state_ph = tf.placeholder(tf.float32, [None, state_size], name="state_ph")
            self.action_ph = tf.placeholder(tf.float32, [None, action_size], name="action_ph")
            self.advantage_ph = tf.placeholder(tf.float32, [None, 1], name="action_ph")

            with tf.variable_scope("pi"):
                self.pi, self.mean_action = self._create_model(trainable=True)
                self.pi_vars = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope="Policy/pi")

                self.sample_op = self.pi.sample()

            with tf.variable_scope("old_pi"):
                self.old_pi, _ = self._create_model(trainable=False)
                self.old_pi_vars = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope="Policy/old_pi")

            with tf.variable_scope("loss"):
                prob_ratio = self.pi.prob(self.action_ph) / self.old_pi.prob(self.action_ph)
                surrogate = prob_ratio * self.advantage_ph

                clipped_surrogate = tf.minimum(surrogate,  tf.clip_by_value(prob_ratio, 1.-epsilon, 1.+epsilon)*self.advantage_ph)

                self.pi_entropy = self.pi.entropy()

                tf.summary.scalar("entropy", tf.reduce_mean(self.pi_entropy))

                self.loss = -tf.reduce_mean(clipped_surrogate + alpha_entropy * self.pi_entropy)

                tf.summary.scalar("objective", self.loss)

            with tf.variable_scope("training"):
                self.gradients = tf.gradients(self.loss, self.pi_vars)
                #self.gradients = [tf.clip_by_value(g, -1000, 1000) for g in self.gradients]
                #self.gradients, _ = tf.clip_by_global_norm(self.gradients, GRADIENT_NORM)
                grads = zip(self.gradients, self.pi_vars)
                self.optimize = tf.train.AdamOptimizer(lr).apply_gradients(grads)

                [tf.summary.histogram(v.name, g) for g, v in grads]

            with tf.variable_scope("update_old_policy"):
                self.update_oldpi_op = [oldp.assign(p) for p, oldp in zip(self.pi_vars, self.old_pi_vars)]

            self.summary_op = tf.summary.merge_all(scope="Policy")

def _create_model(self, trainable):
    layer_names = ["l1", "l2", "l3", "l4"]

    l1 = tf.layers.Dense(32, activation="relu", name=layer_names[0], trainable=trainable, kernel_initializer = tf.initializers.he_normal(),)(self.state_ph)
    l2 = tf.layers.Dense(64, activation="relu", name=layer_names[1], trainable=trainable, kernel_initializer = tf.initializers.he_normal(),)(l1)
    l3 = tf.layers.Dense(32, activation="relu", name=layer_names[2], trainable=trainable, activity_regularizer= tf.contrib.layers.l2_regularizer(scale=0.001),kernel_initializer = tf.initializers.he_normal(),)(l2)
    mu = tf.layers.Dense(self.action_size, activation="tanh", name=layer_names[3], trainable=trainable, kernel_initializer = tf.initializers.he_normal(),)(l3)

    log_sigma = tf.Variable(initial_value=tf.fill((self.action_size,), 0.), trainable=trainable)

    distribution = tfp.distributions.MultivariateNormalDiag(loc=mu, scale_diag=tf.exp(log_sigma))

    tf.summary.histogram("log_sigma", log_sigma)
    tf.summary.histogram("mu", mu)

    for name in layer_names:
        with tf.variable_scope(name, reuse=True):
            tf.summary.histogram("kernel", tf.get_variable("kernel"))
            tf.summary.histogram("bias", tf.get_variable("bias"))

    return distribution, mu

def sample_action(self, state):
    return self.sess.run([self.mean_action, self.sample_op], feed_dict={
        self.state_ph: state
    })

def train(self, states, actions, advantages):
    _, summaries =self.sess.run([self.optimize, self.summary_op], feed_dict={
        self.state_ph:states,
        self.action_ph: actions,
        self.advantage_ph: advantages
    })
    return summaries

def update_old_pi(self):
    self.sess.run([self.update_oldpi_op])

Результат профилирования: