<jupyter_start><jupyter_text>Conv model<jupyter_code>from tensorflow.keras.backend import image_data_format
image_data_format()
learner = Learner(Config())
# using channels last format in the env as well
plt.imshow(learner.env.reset() @ np.random.rand(11, 3))
plt.show()
Xo_train.shape, Xa_train.shape
class Sampling(tf.keras.layers.Layer):
    """Uses (z_mean, z_log_var) to sample z, the vector encoding a digit."""

    def call(self, inputs):
        z_mean, z_log_var = inputs
        batch = tf.shape(z_mean)[0]
        dim = tf.shape(z_mean)[1]
        epsilon = tf.keras.backend.random_normal(shape=(batch, dim))
        return z_mean + tf.exp(0.5 * z_log_var) * epsilon
np.random.shuffle(Xo_train)
from tensorflow.keras.layers import Input, Concatenate, Conv2D, Flatten, Dense, ReLU, MaxPooling2D, UpSampling2D, Conv2DTranspose, Reshape
from tensorflow.keras.models import Model
from tensorflow.keras.activations import tanh

# Define two input layers
image_input = Input((12, 7, 11))
vector_input = Input((4,))

# Convolution + Flatten for the image
conv_layer = image_input

conv_layer = Conv2D(512, (3, 3), padding='valid')(conv_layer)
conv_layer = MaxPooling2D()(conv_layer)
print(conv_layer.shape)
conv_layer = ReLU()(conv_layer)
# conv_layer = Conv2D(64, (5, 5), padding='same')(conv_layer)

# print(conv_layer.shape)

# conv_layer = Conv2D(1024, (2, 2))(conv_layer)
# conv_layer = MaxPooling2D((2, 2), padding='valid')(conv_layer)
# conv_layer = ReLU()(conv_layer)

flat_layer = Flatten()(conv_layer)

print(conv_layer.shape, flat_layer.shape)

# Concatenate the convolutional features and the vector input
# concat_layer = Concatenate()([vector_input, flat_layer])
concat_layer = flat_layer
middle = Dense(1024)(concat_layer)
middle = ReLU()(middle)
mean = Dense(2688)(middle)
logvar = Dense(2688)(middle)

middle = Sampling()([mean, logvar])

upsample = Reshape((12, 7, 32))(middle)

# upsample = UpSampling2D()(upsample)
# upsample = Conv2DTranspose(512, (2, 2))(upsample)
# upsample = ReLU()(upsample)

# print(upsample.shape)

# upsample = UpSampling2D((2, 2))(upsample)
# print(upsample.shape)
upsample = Conv2DTranspose(11, (3, 3), padding='same')(upsample)
# upsample = ReLU()(upsample)

# print(upsample.shape)

# define a model with a list of two inputs
model = Model(inputs=[image_input], outputs=upsample)

print(upsample.shape, image_input.shape)
assert model.predict([Xo_train[:10], Xa_train[:10]]).shape[1:] == yo_train.shape[1:]
model.compile('adam', 'mse')
# fitting on inputs...
model.fit([Xo_train, Xa_train], Xo_train, epochs=100, batch_size=128)
obs1 = model.predict([Xo_train[0:1], Xa_train[0:1]])
# using channels last format in the env as well
plt.imshow(obs1[0, :, :, :] @ np.random.rand(11, 3))
plt.show()<jupyter_output>WARNING:matplotlib.image:Clipping input data to the valid range for imshow with RGB data ([0..1] for floats or [0..255] for integers).<jupyter_text>FC model -- doesn't work well<jupyter_code>pickle.dump([Xtrain, ytrain, Xtest, ytest], open('dataset.pkl', 'wb'))
Xtrain, ytrain, Xtest, ytest = pickle.load(open('dataset.pkl', 'rb'))
gin.bind_parameter('Config.env_steps', 10)
gin.bind_parameter('Config.disable_cuda', True)
learner = Learner(Config())
env = learner.env
h, w, c = env.engine._observation.shape
idxes = list(range(len(Xtrain)))
np.random.shuffle(idxes)
# forgot to shuffle the dataset...
Xtrain = Xtrain[idxes]
ytrain = ytrain[idxes]
!du -s --si dataset.pkl
model = tf.keras.Sequential([
    tf.keras.layers.Dense(64, input_shape=(Xtrain.shape[1],), activation='tanh', kernel_initializer='glorot_normal'),
    tf.keras.layers.Dense(64, activation='tanh', kernel_initializer='glorot_normal'),
    tf.keras.layers.Dense(64, activation='tanh', kernel_initializer='glorot_normal'),
    tf.keras.layers.Dense(2048, activation='tanh', kernel_initializer='glorot_normal'),
    tf.keras.layers.Dense(2048, activation='tanh', kernel_initializer='glorot_normal'),
    tf.keras.layers.Dense(ytrain.shape[1], activation='sigmoid', kernel_initializer='glorot_normal')
])
model.compile('adam', 'binary_crossentropy', metrics=['mse', 'mae'])
model.fit(Xtrain, ytrain, validation_data=(Xtest, ytest), epochs=100)<jupyter_output>Train on 91632 samples, validate on 9184 samples
Epoch 1/100
91632/91632 [==============================] - 49s 533us/sample - loss: 0.0187 - mean_squared_error: 0.0047 - mean_absolute_error: 0.0097 - val_loss: 0.0074 - val_mean_squared_error: 0.0019 - val_mean_absolute_error: 0.0039
Epoch 2/100
91632/91632 [==============================] - 50s 543us/sample - loss: 0.0041 - mean_squared_error: 9.6965e-04 - mean_absolute_error: 0.0021 - val_loss: 0.0043 - val_mean_squared_error: 0.0010 - val_mean_absolute_error: 0.0020
Epoch 3/100
91632/91632 [==============================] - 43s 465us/sample - loss: 0.0033 - mean_squared_error: 7.8184e-04 - mean_absolute_error: 0.0015 - val_loss: 0.0039 - val_mean_squared_error: 9.1098e-04 - val_mean_absolute_error: 0.0017
Epoch 4/100
91632/91632 [==============================] - 46s 504us/sample - loss: 0.0031 - mean_squared_error: 7.3046e-04 - mean_absolute_error: 0.0014 - val_loss: 0.0039 - val_mean_squared_error: 8.9386e-04 - val_mean_absolute_e[...]<jupyter_text>Attention layer<jupyter_code>Xtrain, ytrain = get_xy(100000)
Xtest, ytest = get_xy(10000)
pickle.dump([Xtrain, ytrain, Xtest, ytest], open('dataset.pkl', 'wb'))
from tqdm.auto import tqdm
import numpy as np
import tensorflow as tf
from causal_util.collect_data import EnvDataCollector
from matplotlib import pyplot as plt
import os
os.environ['TF_FORCE_GPU_ALLOW_GROWTH'] = "true"
os.environ['CUDA_VISIBLE_DEVICES'] = "-1"

import sys
from sparse_causal_model_learner_rl.learner import Learner
import gin
from sparse_causal_model_learner_rl.sacred_gin_tune.sacred_wrapper import load_config_files
from sparse_causal_model_learner_rl.config.config import Config
from sparse_causal_model_learner_rl.trainable.decoder import IdentityDecoder
import pickle
%matplotlib inline

load_config_files(['../keychest/config/5x5.gin', '../sparse_causal_model_learner_rl/configs/common.gin'])
Xtrain, ytrain, Xtest, ytest = pickle.load(open('dataset.pkl', 'rb'))
Xtrain.shape
gin.bind_parameter('Config.env_steps', 10)
learner = Learner(Config())
env = learner.env
h, w, c = env.engine._observation.shape
idxes = list(range(len(Xtrain)))
np.random.shuffle(idxes)
# forgot to shuffle the dataset...
Xtrain = Xtrain[idxes]
ytrain = ytrain[idxes]
from tensorflow.keras.layers import Input, Concatenate, Flatten, Dense, ReLU, Reshape, MultiHeadAttention
from tensorflow.keras.models import Model
from tensorflow.keras.activations import tanh

inp = Input((928,))

x = inp
# x = Dense(20, activation='sidmoig')(x)
x = Reshape((928, 1))(x)
print(x.shape)
x, aw = MultiHeadAttention(num_heads=1, key_dim=2, attention_axes=(1, ))(x, x, return_attention_scores=True)
# x, aw = MultiHeadAttention(num_heads=5, key_dim=10, attention_axes=(1, ))(x, x, return_attention_scores=True)

# x = MultiHeadAttention(num_heads=50, key_dim=50, attention_axes=(1,))(x, x)
print(aw.shape, x.shape)
x = Flatten()(x)
x = Dense(924, activation=None)(x)

# define a model with a list of two inputs
model = Model(inputs=[inp], outputs=x)

print(x.shape, inp.shape)
model.compile('adam', 'mse', metrics=['mse', 'mae'])
model.fit(Xtrain, ytrain, validation_data=(Xtest, ytest), epochs=100, batch_size=64)
import tensorflow as tf
tf.test.is_gpu_available()
tf.config.list_physical_devices('GPU')<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1565589099652685, 'fraction_numerical': 0.05009350788137858, 'mean_word_length': 6.318670576735093, 'pattern_counts': {'":': 0, '<': 11, '<?xml version=': 0, '>': 11, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '3', 'dir': 'jupyter-structured-clean-dedup', 'id': '277327', 'license': 'CC0-1.0', 'n_tokens_mistral': 3052, 'n_tokens_neox': 2753, 'n_words': 563, 'path': 'debug/keychestenv-fit-test-keras.ipynb', 'repo_name': 'sergeivolodin/causality-disentanglement-rl', 'size': '7486'}