Как собрать все переменные в виде списка в TensorFlow, сгруппировав их как функцию

Я пытаюсь воспроизвести архитектуру сети cGAN, представленную в недавней статье deep video portrait (2018, Stanford)

Я определил Generator как T(x) согласно обозначениям в статье.

А T(x) относится к вышеупомянутым блокам операций, таким как conv_down(), conv_upsample(), biLinearDown() и finalTanH().

Я указал их область действия с помощью синтаксиса ‘with tf.variable_scope()’.

Пока я составлял функцию потерь и оптимизатор, я обнаружил, что мне нужно собрать все переменные, связанные с генератором, вместе, поскольку мы собираемся обучать с помощью двух разных оптимизаторов, один для дискриминатора и один для генератора.

Дискриминатор находится под контролем моего коллеги, поэтому это не моя забота, и я оставляю его в виде псевдокода.

Тем не менее, я хотел бы сделать список переменных, определенных в T(x) в моем коде.

Как я могу это сделать? Любая помощь?

import tensorflow as tf
import numpy as np

# гиперпараметры
learning_rate = 0.0002
epochs = 250
batch_size = 16
N_w = 11 # количество связанных кадров
channels = 9*N_w
drop_out = [0.5, 0.5, 0.5, 0, 0, 0, 0, 0]
lambda_ = 100 # для взвешивания T_loss

tf.reset_default_graph()

with tf.Graph().as_default():

    def conv_down(x, N, count): #Conv [4x4, str_2] > Batch_Normalization > Leaky_ReLU
        with tf.variable_scope("conv_down_{}_count{}".format(N, count)): #N == глубина тензора
            x = tf.layers.conv2d(x, N, kernel_size=4, strides=2, padding='same', kernel_initializer=tf.truncated_normal_initializer(stddev=np.sqrt(0.2)))
            x = tf.contrib.layers.batch_norm(x)
            x = tf.nn.leaky_relu(x) # для conv_down, реализовать leakyReLU
        return x

    def conv_upsample(x, N, drop_rate, count):
        with tf.variable_scope("conv_upsamp_{}_count{}".format(N,count)):
            # увеличение
            with tf.variable_scope("conv_up_count{}".format(count)):
                x = tf.layers.conv2d_transpose(x, N, kernel_size=4, strides=2, padding='same', kernel_initializer=tf.truncated_normal_initializer(stddev=np.sqrt(0.2)))
                x = tf.contrib.layers.batch_norm(x)
                with tf.variable_scope("convdrop_{}".format(count)):
                    if drop_rate is not 0:
                        x = tf.nn.dropout(x, keep_prob=drop_rate)
                x = tf.nn.relu(x)

            # уточнение1
            with tf.variable_scope("refine1"):
                x = tf.layers.conv2d(x, N, kernel_size=3, strides=1, padding='same', kernel_initializer=tf.truncated_normal_initializer(stddev=np.sqrt(0.2)))
                x = tf.contrib.layers.batch_norm(x)
                with tf.variable_scope("rf1drop_out_{}".format(count)):
                    if drop_rate is not 0:
                        x = tf.nn.dropout(x, keep_prob=drop_rate)
                x = tf.nn.relu(x)

            # уточнение2
            with tf.variable_scope("refine2"):
                x = tf.layers.conv2d(x, N, kernel_size=3, strides=1, padding='same', kernel_initializer=tf.truncated_normal_initializer(stddev=np.sqrt(0.2)))
                x = tf.contrib.layers.batch_norm(x)
                with tf.variable_scope("rf2drop_out{}".format(count)):
                    if drop_rate is not 0:
                        x = tf.nn.dropout(x, keep_prob=drop_rate)
                x = tf.nn.relu(x)

        return x

    def biLinearDown(x, N):
        return tf.image.resize_images(x, [N, N])

    def finalTanH(x):
        with tf.variable_scope("tanh"):
            x = tf.nn.tanh(x)
        return x

    def T(x):
        # структура выходного канала
        down_channel_output = [64, 128, 256, 512, 512, 512, 512, 512]
        up_channel_output= [512, 512, 512, 512, 256, 128, 64, 3]
        biLinearDown_output= [32, 64, 128] # для пропускного соединения

        # понижающая частота
        conv1 = conv_down(x, down_channel_output[0], 1)
        conv2 = conv_down(conv1, down_channel_output[1], 2)
        conv3 = conv_down(conv2, down_channel_output[2], 3)
        conv4 = conv_down(conv3, down_channel_output[3], 4)
        conv5 = conv_down(conv4, down_channel_output[4], 5)
        conv6 = conv_down(conv5, down_channel_output[5], 6)
        conv7 = conv_down(conv6, down_channel_output[6], 7)
        conv8 = conv_down(conv7, down_channel_output[7], 8)

        # увеличение частоты
        dconv1 = conv_upsample(conv8, up_channel_output[0], drop_out[0], 1)
        dconv2 = conv_upsample(dconv1, up_channel_output[1], drop_out[1], 2)
        dconv3 = conv_upsample(dconv2, up_channel_output[2], drop_out[2], 3)
        dconv4 = conv_upsample(dconv3, up_channel_output[3], drop_out[3], 4)
        dconv5 = conv_upsample(dconv4, up_channel_output[4], drop_out[4], 5)
        dconv6 = conv_upsample(tf.concat([dconv5, biLinearDown(x, biLinearDown_output[0])], axis=3), up_channel_output[5], drop_out[5], 6)
        dconv7 = conv_upsample(tf.concat([dconv6, biLinearDown(x, biLinearDown_output[1])], axis=3), up_channel_output[6], drop_out[6], 7)
        dconv8 = conv_upsample(tf.concat([dconv7, biLinearDown(x, biLinearDown_output[2])], axis=3), up_channel_output[7], drop_out[7], 8)

        # финальный tanh
        T_x = finalTanH(dconv8)

        return T_x

    # входной тензор x : для подачи как Fake
    x = tf.placeholder(tf.float32, [batch_size, 256, 256, channels]) # batch_size x Высота x Ширина x N_w

    # сгенерированный тензор T(x)
    T_x = T(x)

    # целевой тензор Y : для подачи как Real
    Y = tf.placeholder(tf.float32, [batch_size, 256, 256, 3]) # просто кадр видео

    # определяем псевдо дискриминатор
    def D(x, to_be_discriminated): # истина может быть либо T(x), либо GroundTruth с формой [256 x 256 x 3]
        sheudo_prob = np.float32(np.random.uniform(low=0., high=1.))
        return sheudo_prob

    theta_D = [] # tf.Variables для Дискриминатора

    # Дискриминированный Результат
    D_real = D(Y)
    D_fake = D(T_x)

    # Определяем функцию потерь
    E_cGAN = tf.reduce_mean(tf.log(D_real)+ tf.log(1. - D_fake))
    E_l1 = tf.reduce_mean(tf.norm((Y-T_x)))
    Loss = EcGAN + lambda_*E_l1

    # Оптимизатор
    D_solver = tf.train.AdamOptimizer().minimize(-Loss, var_list=theta_D)   # Обновляем только параметры D(X), поэтому var_list = theta_D
    T_solver = tf.train.AdamOptimizer().minimize(Loss, var_list=theta_T)   # Обновляем только параметры G(X), поэтому var_list = theta_T


    ####ТЕСТ####
    # определяем псевдо ввод для тестирования
    sheudo_x = np.float32(np.random.uniform(low=-1., high=1., size=[16, 256, 256, 99]))
    sheudo_Y = np.float32(np.random.uniform(low=-1., high=1., size=[16, 256, 256, 3]))


    ####Запуск####

    init_g = tf.global_variables_initializer()
    with tf.Session() as sess:
        sess.run(init_g)
        sess.run(output_tensor,  feed_dict={x: sheudo_input Y: sheudo_Y})

В GANs вам необходимо тренировать некоторые параметры, зафиксировать их и тренировать другие, и эта операция может выполняться многократно. Вы можете выполнить следующую последовательность операций.

Укажите все переменные, связанные с генератором, внутри их соответствующей области видимости переменной, а затем получите доступ к ним, используя tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope="my_scope"). После этого во время обучения вы можете передать эти переменные как обучаемые параметры вашему оптимизатору, установив var_list в методе minimize. Вы также можете ознакомиться с этим.

Если вы хотите получить все обучаемые переменные, вы можете собрать их в список, используя метод tf.trainable_variables.

Возможно, стоит взглянуть на этот пост для других аспектов зафиксирования переменных.

Вы также можете посмотреть на реализацию сети адверсаров Hvass-Labs. Также взгляните здесь.