我可以在model_fn，Estimator，Tensorflow中使用python循环吗？

Question

我对tensorflow估计器非常陌生，我试图“用ConvNet处理每一帧视频，总结重建损失，然后优化参数”。

因此，我想知道我是否可以用为循环编写model_fn作为估计器，这样我就可以处理每一帧视频，然后一起进行优化。

谢谢

我附上两段自己的实现，它们都能工作。即使我将model_fn cnn_model嵌入到model_fn中(以前是不允许的)，估计器的定义似乎也是允许循环的。

Vanilla执行：

import os, sys
import tensorflow as tf

# Read in video dataset in [batch, frames, height, width]
raw_data = np.load(Data_root)
dataset = tf.data.Dataset.from_tensor_slices((raw_data))
dataset = dataset.batch(BATCH_SIZE)
iterator = dataset.make_initializable_iterator()
one_element = iterator.get_next()

# Set up placeholder for each frame
frame = tf.placeholder('float32', [BATCH_SIZE, IMAGE_HEIGHT, IAMGE_WIDTH])
label = tf.placeholder('float32', [BATCH_SIZE, IMAGE_HEIGHT, IAMGE_WIDTH])

# Define network
with tf.name_scope("network"):

    with tf.name_scope("Encoder"):
        conv1 = tf.layers.conv2d(frames, 32, [3,3], strides=2, padding='same', activation=tf.nn.relu)
        conv2 = tf.layers.conv2d(conv1, 64, [3,3], strides=2, padding='same', activation=tf.nn.relu)

    with tf.name_scope("Repeat_Layer"):
        latent = tf.layers.conv2d(conv2, 64, [3,3], strides=2, padding='same', activation=tf.nn.relu)

    with tf.name_scope("Decoder"):
        conv3 = tf.layers.conv2d_transpose(latent, 32, [3, 3], strides=2, padding='same', activation=tf.nn.relu)
        conv4 = tf.layers.conv2d_transpose(conv3, 1, [3, 3], strides=2, padding='same', activation=tf.nn.relu)

prediction = tf.identity(conv4, name='prediciton')

# Define loss
loss_mse = tf.losses.mean_squared_error(frame, prediction)

# Define optim
optimizer = tf.train.RMSPropOptimizer(0.001).minimize(loss_total)

# Init
init_op = tf.group(tf.global_variables_initializer(), tf.local_variables_initializer())

with tf.Session as sess:

    sess.run(init_op)
    coord = tf.train.Coordinator()
    threads = tf.train.start_queue_runners(sess=sess, coord=coord)

    # Assume train_batch with shape [batch, 5, height, width]
    train_batch, label_batch = sess.run([one_element])
    loss_total = 0.0

    for i in range(5):

        feed_dict = {frame:train_batch[:, i, :, :], label:label_batch[0][:, i, :, :]}
        loss = sess.run([loss_mse], feed_dict=feed_dict)
        loss_total += loss

    feed_dict = {loss_total:loss_total}
    _ = sess.run([optimizer], feed_dict=feed_dict)

print("Optimization is Finished!")
coord.request_stop()
coord.join(threads)
sess.close()

估计数执行情况：

import os, sys
import tensorflow as tf
import numpy as np

def cnn_model(input_feature):

        with tf.name_scope("Encoder"):
            conv1 = tf.layers.conv2d(frames, 32, [3,3], strides=2, padding='same', activation=tf.nn.relu)
            conv2 = tf.layers.conv2d(conv1, 64, [3,3], strides=2, padding='same', activation=tf.nn.relu)

        with tf.name_scope("Repeat_Layer"):
            latent = tf.layers.conv2d(conv2, 64, [3,3], strides=2, padding='same', activation=tf.nn.relu)

        with tf.name_scope("Decoder"):
            conv3 = tf.layers.conv2d_transpose(latent, 32, [3, 3], strides=2, padding='same', activation=tf.nn.relu)
            conv4 = tf.layers.conv2d_transpose(conv3, 1, [3, 3], strides=2, padding='same', activation=tf.nn.relu)

    return conv4

def model_fn(features, labels, mode):

    # Assume each video contains five frames
    input_feature = tf.reshape(features, [batch, 5, height, width])
    loss_total = 0.0

    for i in range(5):

        input_layer = input_feature[:, i, :, :]

        prediction = cnn_model(input_layer)

        loss_mse = tf.losses.mean_squared_error(labels=labels, predictions=prediction)

        loss_total += loss_mse

    if mode == tf.estimator.ModeKeys.TRAIN:

        optimizer = tf.train.GradientDescentOptimizer(learning_rate=0.001)
        train_op = optimizer.minimize(
            loss=loss_mse,
            global_step=tf.train.get_global_step()
        )

    return tf.estimator.EstimatorSpec(mode=mode, loss=loss_mse, train_op=train_op)

def main(unused_argv):

    # Load video data [batch, 5, height, width]
    data_path = '/xxx/train.npy'
    train_data = np.load(data_path)

    # Set up Estimator
    AutoEncoder = tf.estimator.Estimator(
        model_fn=model_fn, model_dir=None
    )

    # Set up input_fn pipeline
    train_input_fn = tf.estimator.inputs.numpy_input_fn(
        x={"x":train_data},
        y=train_data,
        batch_size=10,
        num_epochs=100,
        shuffle=True
    )

    # Start train
    AutoEncoder.train(
        input_fn=train_input_fn,
        steps=15000,
        hooks=hooks
    )

if __name__ == "__main__":

    with tf.device("/gpu:0"):
        tf.app.run()

Answer 1

我会在这里回答我自己的问题。

答案是，是，。在估值器框架下，我们可以在model_fn中使用python循环，如上面的第二段代码所示。

如果运行第二个sinppet，考虑到Estimator本身将生成Tensorboard日志记录，我们可以很容易地通过"tensorboard --logdir='path_to_model'“签出图结构。您将看到AutoEncoder模块运行了5次，损失也增加了5倍，这证明了我的猜测。

在某种程度上，这是一个重要的问题。利用for循环，我们可以输入任意类型的序列数据，处理每一段数据，然后一起优化模型。例如，我可以处理每一帧的视频，计算重建损失，然后返回-支撑整体损失。