openai / iaf

Code for reproducing key results in the paper "Improving Variational Inference with Inverse Autoregressive Flow"
https://arxiv.org/abs/1606.04934
MIT License
518 stars 133 forks source link

Merge pull request #12 from openai/revert-8-tf_utils_bugfix #14

Open ak-kat opened 6 years ago

ak-kat commented 6 years ago

class IAFLayer(object): def init(self, hps, mode, downsample): self.hps = hps self.mode = mode self.downsample = downsample

def up(self, input, **_):
    hps = self.hps
    h_size = hps.h_size
    z_size = hps.z_size
    stride = [2, 2] if self.downsample else [1, 1]

    with arg_scope([conv2d]):
        x = tf.nn.elu(input)
        x = conv2d("up_conv1", x, 2 * z_size + 2 * h_size, stride=stride)
        self.qz_mean, self.qz_logsd, self.up_context, h = split(x, 1, [z_size, z_size, h_size, h_size])

        h = tf.nn.elu(h)

h = conv2d("up_conv3", h, h_size)

--------------------------------------------------- Running on a single gpu ---------------------------------------------- conv2d function execution is eating up huge memory. Initial two execution of up function are executing fine when called from forward sub_layer.up(h).

def _forward(self, x, gpu): hps = self.hps

    x = tf.to_float(x)
    x = tf.clip_by_value((x + 0.5) / 256.0, 0.0, 1.0) - 0.5

    # Input images are repeated k times on the input.
    # This is used for Importance Sampling loss (k is number of samples).
    data_size = hps.batch_size * hps.k
    x = repeat(x, hps.k)

    orig_x = x
    h_size = hps.h_size

    with arg_scope([conv2d, deconv2d], init=(self.mode == "init")):
        layers = []
        for i in range(hps.depth):
            layers.append([])
            for j in range(hps.num_blocks):
                downsample = (i > 0) and (j == 0)
                layers[-1].append(IAFLayer(hps, self.mode, downsample))

        h = conv2d("x_enc", x, h_size, [5, 5], [2, 2])  # -> [16, 16]
        for i, layer in enumerate(layers):
            for j, sub_layer in enumerate(layer):
                with tf.variable_scope("IAF_%d_%d" % (i, j)):

h = sub_layer.up(h)

* but in further iterartion it ate up memory of 100GB and continuously increasing. I have forcefully stopped the execution. Please provide the solution for the increased memory issue for conv2d function.*****

Thanks in advance.