Python tensorflow.contrib.slim 模块，stack() 实例源码

def fc_net(inp, layers, out_layers, scope, lamba=1e-3, activation=tf.nn.relu, reuse=None,
           weights_initializer=initializers.xavier_initializer(uniform=False)):
    with slim.arg_scope([slim.fully_connected],
                        activation_fn=activation,
                        normalizer_fn=None,
                        weights_initializer=weights_initializer,
                        reuse=reuse,
                        weights_regularizer=slim.l2_regularizer(lamba)):

        if layers:
            h = slim.stack(inp, slim.fully_connected, layers, scope=scope)
            if not out_layers:
                return h
        else:
            h = inp
        outputs = []
        for i, (outdim, activation) in enumerate(out_layers):
            o1 = slim.fully_connected(h, outdim, activation_fn=activation, scope=scope + '_{}'.format(i + 1))
            outputs.append(o1)
        return outputs if len(outputs) > 1 else outputs[0]

def encoder(x): 
    # Variational posterior q(y|x), i.e. the encoder (shape=(batch_size, 200))
    net = slim.stack(x, 
                     slim.fully_connected, 
                     [512, 256])

    # Unnormalized logits for number of classes (N) seperate K-categorical distributions
    logits_y = tf.reshape(slim.fully_connected(net, 
                                               FLAGS.num_classes*FLAGS.num_cat_dists, 
                                               activation_fn=None), 
                          [-1, FLAGS.num_cat_dists])

    q_y = tf.nn.softmax(logits_y)
    log_q_y = tf.log(q_y + 1e-20)

    return logits_y, q_y, log_q_y

def decoder(tau, logits_y):
    y = tf.reshape(gumbel_softmax(logits_y, tau, hard=False), 
                   [-1, FLAGS.num_cat_dists, FLAGS.num_classes])

    # Generative model p(x|y), i.e. the decoder (shape=(batch_size, 200))
    net = slim.stack(slim.flatten(y), 
                     slim.fully_connected, 
                     [256, 512])

    logits_x = slim.fully_connected(net, 
                                    784, 
                                    activation_fn=None)

    # (shape=(batch_size, 784))
    p_x = bernoulli(logits=logits_x)

    return p_x

def get_model(self, inputs):
        # get the normalizer function and parameters
        normalizer_fn, normalizer_params = _get_normalizer(self.is_training,
                                                           self.use_bn,
                                                           self.use_ln)
        winit = tf.contrib.layers.xavier_initializer()
        binit = tf.constant_initializer(0)

        with tf.variable_scope(self.scope):
            with slim.arg_scope([slim.fully_connected],
                                activation_fn=self.activation,
                                weights_initializer=winit,
                                biases_initializer=binit,
                                normalizer_fn=normalizer_fn,
                                normalizer_params=normalizer_params):
                    layers = slim.stack(inputs, slim.fully_connected,
                                        self.sizes, scope="layer")

            output_size = self.latent_size * self.double_features
            return slim.fully_connected(layers, output_size,
                                        activation_fn=None,
                                        normalizer_fn=None,
                                        weights_initializer=winit,
                                        biases_initializer=binit,
                                        scope='projection')

def actor_network():
    '''
    Actor network, including policy gradient equation and optimizer
    '''
    with tf.variable_scope('policy'):
        # Inputs
        state = tf.placeholder('float', [None, OBS_WIDTH])  # batch_size x obs_width
        actions = tf.placeholder('float', [None, NUM_ACTIONS])  # batch_size x num_actions
        advantages = tf.placeholder('float', [None, 1])  # batch_size x 1

        # 3-layer fully-connected neural network
        mlp_out = slim.stack(state, slim.fully_connected, [6, NUM_ACTIONS], weights_regularizer=slim.l2_regularizer(scale=A_REG_SCALE))

        # Network output
        probabilities = tf.nn.softmax(mlp_out)

        good_probabilities = tf.reduce_sum(tf.mul(probabilities, actions), reduction_indices=[1])
        eligibility = tf.log(good_probabilities) * advantages

        # Loss & optimizer
        data_loss = -tf.reduce_sum(eligibility)
        reg_losses = slim.losses.get_regularization_losses(scope='policy')
        reg_loss = tf.reduce_sum(reg_losses)
        total_loss = data_loss + reg_loss

        optimizer = tf.train.AdamOptimizer(ACTOR_LR).minimize(total_loss)

        return probabilities, state, actions, advantages, optimizer

def critic_network():
    '''
    Critic network, including loss and optimizer
    '''
    with tf.variable_scope('value'):
        # Inputs
        state = tf.placeholder('float', [None, OBS_WIDTH])  # batch_size x obs_width
        newvals = tf.placeholder('float', [None, 1])  # batch_size x 1

        # 4-layer fully-connected neural network
        calculated = slim.stack(state, slim.fully_connected, [6, 6, 1], weights_regularizer=slim.l2_regularizer(scale=C_REG_SCALE))

        # Error value
        diffs = calculated - newvals

        # Loss & optimizer
        data_loss = tf.nn.l2_loss(diffs)
        reg_losses = slim.losses.get_regularization_losses(scope='value')
        reg_loss = tf.reduce_sum(reg_losses)
        total_loss = data_loss + reg_loss

        optimizer = tf.train.AdamOptimizer(CRITIC_LR).minimize(total_loss)

        return calculated, state, newvals, optimizer, total_loss

########################################
# Training and inference processes
########################################

def encoder(self, x):
        """ Convolutional variational encoder to encode image into a low-dimensional latent code
        If config.conv == False it is a MLP VAE. If config.use_vae == False, it is a normal encoder
        :param x: sequence of images
        :return: a, a_mu, a_var
        """
        with tf.variable_scope('vae/encoder'):
            if self.config.conv:
                x_flat_conv = tf.reshape(x, (-1, self.d1, self.d2, 1))
                enc_hidden = slim.stack(x_flat_conv,
                                        slim.conv2d,
                                        self.num_filters,
                                        kernel_size=self.config.filter_size,
                                        stride=2,
                                        activation_fn=self.activation_fn,
                                        padding='SAME')
                enc_flat = slim.flatten(enc_hidden)
                self.enc_shape = enc_hidden.get_shape().as_list()[1:]

            else:
                x_flat = tf.reshape(x, (-1, self.d1 * self.d2))
                enc_flat = slim.repeat(x_flat, self.config.num_layers, slim.fully_connected,
                                       self.config.vae_num_units, self.activation_fn)

            a_mu = slim.fully_connected(enc_flat, self.config.dim_a, activation_fn=None)
            if self.config.use_vae:
                a_var = slim.fully_connected(enc_flat, self.config.dim_a, activation_fn=tf.nn.sigmoid)
                a_var = self.config.noise_emission * a_var
                a = simple_sample(a_mu, a_var)
            else:
                a_var = tf.constant(1., dtype=tf.float32, shape=())
                a = a_mu
            a_seq = tf.reshape(a, tf.stack((-1, self.ph_steps, self.config.dim_a)))
        return a_seq, a_mu, a_var

def decoder(self, a_seq):
        """ Convolutional variational decoder to decode latent code to image reconstruction
        If config.conv == False it is a MLP VAE. If config.use_vae == False it is a normal decoder
        :param a_seq: latent code
        :return: x_hat, x_mu, x_var
        """
        # Create decoder
        if self.config.out_distr == 'bernoulli':
            activation_x_mu = tf.nn.sigmoid
        else:
            activation_x_mu = None

        with tf.variable_scope('vae/decoder'):
            a = tf.reshape(a_seq, (-1, self.config.dim_a))
            if self.config.conv:
                dec_upscale = slim.fully_connected(a, int(np.prod(self.enc_shape)), activation_fn=None)
                dec_upscale = tf.reshape(dec_upscale, [-1] + self.enc_shape)

                dec_hidden = dec_upscale
                for filters in reversed(self.num_filters):
                    dec_hidden = slim.conv2d(dec_hidden, filters * 4, self.config.filter_size,
                                             activation_fn=self.activation_fn)
                    dec_hidden = subpixel_reshape(dec_hidden, 2)
                x_mu = slim.conv2d(dec_hidden, 1, 1, stride=1, activation_fn=activation_x_mu)
                x_var = tf.constant(self.config.noise_pixel_var, dtype=tf.float32, shape=())
            else:
                dec_hidden = slim.repeat(a, self.config.num_layers, slim.fully_connected,
                                         self.config.vae_num_units, self.activation_fn)

                x_mu = slim.fully_connected(dec_hidden, self.d1 * self.d2, activation_fn=activation_x_mu)
                x_mu = tf.reshape(x_mu, (-1, self.d1, self.d2, 1))
                # x_var is not used for bernoulli outputs. Here we fix the output variance of the Gaussian,
                # we could also learn it globally for each pixel (as we did in the pendulum experiment) or through a
                # neural network.
                x_var = tf.constant(self.config.noise_pixel_var, dtype=tf.float32, shape=())

        if self.config.out_distr == 'bernoulli':
            # For bernoulli we show the probabilities
            x_hat = x_mu
        else:
            x_hat = simple_sample(x_mu, x_var)

        return tf.reshape(x_hat, tf.stack((-1, self.ph_steps, self.d1, self.d2))), x_mu, x_var

def build_model(self):
        # Encoder q(a|x)
        a_seq, a_mu, a_var = self.encoder(self.x)
        a_vae = a_seq

        # Initial state for the alpha RNN
        dummy_lstm = BasicLSTMCell(self.config.alpha_units * 2 if self.config.learn_u else self.config.alpha_units)
        state_init_rnn = dummy_lstm.zero_state(self.config.batch_size, tf.float32)

        # Initialize Kalman filter (LGSSM)
        self.kf = KalmanFilter(dim_z=self.config.dim_z,
                               dim_y=self.config.dim_a,
                               dim_u=self.config.dim_u,
                               dim_k=self.config.K,
                               A=self.init_vars['A'],  # state transition function
                               B=self.init_vars['B'],  # control matrix
                               C=self.init_vars['C'],  # Measurement function
                               R=self.init_vars['R'],  # measurement noise
                               Q=self.init_vars['Q'],  # process noise
                               y=a_seq,  # output
                               u=None,
                               mask=self.mask,
                               mu=self.init_vars['mu'],
                               Sigma=self.init_vars['Sigma'],
                               y_0=self.init_vars['a_0'],
                               alpha=self.alpha,
                               state=state_init_rnn
                               )

        # Get smoothed posterior over z
        smooth, A, B, C, alpha_plot = self.kf.smooth()

        # Get filtered posterior, used only for imputation plots
        filter, _, _, C_filter, _ = self.kf.filter()

        # Get a from the prior z (for plotting)
        a_mu_pred = tf.matmul(C, tf.expand_dims(smooth[0], 2), transpose_b=True)
        a_mu_pred_seq = tf.reshape(a_mu_pred, tf.stack((-1, self.ph_steps, self.config.dim_a)))
        if self.config.sample_z:
            a_seq = a_mu_pred_seq

        # Decoder p(x|a)
        x_hat, x_mu, x_var = self.decoder(a_seq)

        # Compute variables for generation from the model (for plotting)
        self.n_steps_gen = self.config.n_steps_gen  # We sample for this many iterations,
        self.out_gen_det = self.kf.sample_generative_tf(smooth, self.n_steps_gen, deterministic=True,
                                                        init_fixed_steps=self.config.t_init_mask)
        self.out_gen = self.kf.sample_generative_tf(smooth, self.n_steps_gen, deterministic=False,
                                                    init_fixed_steps=self.config.t_init_mask)
        self.out_gen_det_impute = self.kf.sample_generative_tf(smooth, self.test_data.timesteps, deterministic=True,
                                                               init_fixed_steps=self.config.t_init_mask)
        self.out_alpha, _, _, _ = self.alpha(self.a_prev, state=state_init_rnn, u=None, init_buffer=True, reuse=True)

        # Collect generated model variables
        self.model_vars = dict(x_hat=x_hat, x_mu=x_mu, x_var=x_var,
                               a_seq=a_seq, a_mu=a_mu, a_var=a_var, a_vae=a_vae,
                               smooth=smooth, A=A, B=B, C=C, alpha_plot=alpha_plot,
                               a_mu_pred_seq=a_mu_pred_seq, filter=filter, C_filter=C_filter)

        return self