Tensorflow while_loop帧连接

时间:2018-06-11 14:42:57

标签: python tensorflow machine-learning openai-gym

我在连接while_loop中的帧时遇到问题。 首先,一个版本可以工作,但在一个简单的for循环中。 这个版本非常慢,这就是我想在tf.while_loop

中运行此代码的原因
import gym
import tensorflow as tf 
import matplotlib.pyplot as plt
import time

env = gym.make("Pong-v0")

def preprocess(frame):
    with tf.variable_scope('frame_process'):
        output_frame = tf.image.rgb_to_grayscale(frame)
        output_frame = tf.image.crop_to_bounding_box(output_frame, 35, 0, 160, 160)
        output_frame = tf.image.resize_images(output_frame,[80,80],method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
        return tf.squeeze(output_frame)

start_time = time.time()
frame = env.reset()
state = preprocess(frame)
stack = tf.stack(4 * [state], axis=2)
step_op = lambda action : env.step(action)[:3]

with tf.Session() as session:

    session.run(tf.global_variables_initializer())

    for i in range(100):
        action = tf.py_func(lambda : env.action_space.sample(),[],[tf.int64])
        frame_, reward, done = tf.py_func(step_op,[action], [tf.uint8, tf.double, tf.bool])
        state_ = preprocess(frame_)
        stack = tf.concat([stack[:,:,1:], tf.expand_dims(session.run(state_),2)], axis=2)

    print("Done in --- %s seconds ---" % (time.time() - start_time))
    stack2 = session.run(stack)

    for x in range(4):
        plt.imshow(stack2[:,:,x], cmap='gray')
        plt.show()
  

注意:在此版本中,我需要在需要时运行session.run(state_)   tf.expand_dims因为如果我没有,我会收到Illegal Instruction!,图片也会损坏。我不是没有原因......

这是我的第二个版本,带有while_loop:

import tensorflow as tf
import gym
import matplotlib.pyplot as plt
import time

env = gym.make("Pong-v0")

def preprocess(frame):
    with tf.variable_scope('frame_process'):
        output_frame = tf.image.rgb_to_grayscale(frame)
        output_frame = tf.image.crop_to_bounding_box(output_frame, 35, 0, 160, 160)
        output_frame = tf.image.resize_images(output_frame,[80,80],method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
        return tf.squeeze(output_frame)

def body(index,stack):
    action = tf.py_func(lambda : env.action_space.sample(),[],[tf.int64])
    frame_, reward, done = tf.py_func(step_op,[action], [tf.uint8, tf.double, tf.bool])
    state_ = preprocess(frame_)
    state_.set_shape((80,80))
    next_stack = tf.concat([stack[:,:,1:], tf.expand_dims(state_,2)], axis=2)
    return tf.add(index, 1), next_stack


start_time = time.time()

frame = env.reset()
state = preprocess(frame)
stack = tf.stack(4 * [state], axis=2)

i = tf.constant(0)
STEPS = tf.constant(100)

while_condition = lambda i, stack: tf.less(i, STEPS)
step_op = lambda action : env.step(action)[:3]
loop_result = tf.while_loop(while_condition, body, (i, stack))

with tf.Session() as session:

    session.run(tf.global_variables_initializer())
    idx, s = session.run(loop_result)

    print("Done in --- %s seconds ---" % (time.time() - start_time))


    for x in range(4):
        plt.imshow(s[:,:,x], cmap='gray')
        plt.show()

当我运行此代码时,我得到一个Illegal Instruction!并且图像不正确。我认为这是因为在扩展之前我无法评估frame_。在第一个示例中执行frame_ tf.expand_dims之前是否可以评估我的tf.while_loop

0 个答案:

没有答案