def main()

in ebm_combine.py [0:0]
86 lines of code
19 McCabe index (conditional complexity)

def main():
    data = np.load(FLAGS.dsprites_path)['imgs']
    l = latents = np.load(FLAGS.dsprites_path)['latents_values']

    np.random.seed(1)
    idx = np.random.permutation(data.shape[0])

    data = data[idx]
    latents = latents[idx]

    config = tf.ConfigProto()
    sess = tf.Session(config=config)

    # Model 1 will be conditioned on size
    model_size = DspritesNet(num_filters=FLAGS.num_filters, cond_size=True)
    weight_size = model_size.construct_weights('context_0')

    # Model 2 will be conditioned on shape
    model_shape = DspritesNet(num_filters=FLAGS.num_filters, cond_shape=True)
    weight_shape = model_shape.construct_weights('context_1')

    # Model 3 will be conditioned on position
    model_pos = DspritesNet(num_filters=FLAGS.num_filters, cond_pos=True)
    weight_pos = model_pos.construct_weights('context_2')

    # Model 4 will be conditioned on rotation
    model_rot = DspritesNet(num_filters=FLAGS.num_filters, cond_rot=True)
    weight_rot = model_rot.construct_weights('context_3')

    sess.run(tf.global_variables_initializer())
    save_path_size = osp.join(FLAGS.logdir, FLAGS.exp_size, 'model_{}'.format(FLAGS.resume_size))

    v_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='context_{}'.format(0))
    v_map = {(v.name.replace('context_{}'.format(0), 'context_0')[:-2]): v for v in v_list}

    if FLAGS.cond_scale:
        saver = tf.train.Saver(v_map)
        saver.restore(sess, save_path_size)

    save_path_shape = osp.join(FLAGS.logdir, FLAGS.exp_shape, 'model_{}'.format(FLAGS.resume_shape))

    v_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='context_{}'.format(1))
    v_map = {(v.name.replace('context_{}'.format(1), 'context_0')[:-2]): v for v in v_list}

    if FLAGS.cond_shape:
        saver = tf.train.Saver(v_map)
        saver.restore(sess, save_path_shape)


    save_path_pos = osp.join(FLAGS.logdir, FLAGS.exp_pos, 'model_{}'.format(FLAGS.resume_pos))
    v_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='context_{}'.format(2))
    v_map = {(v.name.replace('context_{}'.format(2), 'context_0')[:-2]): v for v in v_list}
    saver = tf.train.Saver(v_map)

    if FLAGS.cond_pos:
        saver.restore(sess, save_path_pos)


    save_path_rot = osp.join(FLAGS.logdir, FLAGS.exp_rot, 'model_{}'.format(FLAGS.resume_rot))
    v_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='context_{}'.format(3))
    v_map = {(v.name.replace('context_{}'.format(3), 'context_0')[:-2]): v for v in v_list}
    saver = tf.train.Saver(v_map)

    if FLAGS.cond_rot:
        saver.restore(sess, save_path_rot)

    X_NOISE = tf.placeholder(shape=(None, 64, 64), dtype=tf.float32)
    LABEL_SIZE = tf.placeholder(shape=(None, 1), dtype=tf.float32)
    LABEL_SHAPE = tf.placeholder(shape=(None, 3), dtype=tf.float32)
    LABEL_POS = tf.placeholder(shape=(None, 2), dtype=tf.float32)
    LABEL_ROT = tf.placeholder(shape=(None, 2), dtype=tf.float32)

    x_mod = X_NOISE

    kvs = {}
    kvs['X_NOISE'] = X_NOISE
    kvs['LABEL_SIZE'] = LABEL_SIZE
    kvs['LABEL_SHAPE'] = LABEL_SHAPE
    kvs['LABEL_POS'] = LABEL_POS
    kvs['LABEL_ROT'] = LABEL_ROT
    kvs['model_size'] = model_size
    kvs['model_shape'] = model_shape
    kvs['model_pos'] = model_pos
    kvs['model_rot'] = model_rot
    kvs['weight_size'] = weight_size
    kvs['weight_shape'] = weight_shape
    kvs['weight_pos'] = weight_pos
    kvs['weight_rot'] = weight_rot

    save_exp_dir = osp.join(FLAGS.savedir, '{}_{}_joint'.format(FLAGS.exp_size, FLAGS.exp_shape))
    if not osp.exists(save_exp_dir):
        os.makedirs(save_exp_dir)


    if FLAGS.task == 'conceptcombine':
        conceptcombine(sess, kvs, data, latents, save_exp_dir)
    elif FLAGS.task == 'labeldiscover':
        labeldiscover(sess, kvs, data, latents, save_exp_dir)
    elif FLAGS.task == 'gentest':
        save_exp_dir = osp.join(FLAGS.savedir, '{}_{}_gen'.format(FLAGS.exp_size, FLAGS.exp_pos))
        if not osp.exists(save_exp_dir):
            os.makedirs(save_exp_dir)

        gentest(sess, kvs, data, latents, save_exp_dir)
    elif FLAGS.task == 'genbaseline':
        save_exp_dir = osp.join(FLAGS.savedir, '{}_{}_gen_baseline'.format(FLAGS.exp_size, FLAGS.exp_pos))
        if not osp.exists(save_exp_dir):
            os.makedirs(save_exp_dir)

        if FLAGS.plot_curve:
            mse_losses = []
            for frac in [i/10 for i in range(11)]:
                mse_loss = genbaseline(sess, kvs, data, latents, save_exp_dir, frac=frac)
                mse_losses.append(mse_loss)
            np.save("mse_baseline_comb.npy", mse_losses)
        else:
            genbaseline(sess, kvs, data, latents, save_exp_dir)