<jupyter_start><jupyter_text>---<jupyter_code>input_shape = (10,10,3)
mask_type = 'A'
kernel_size = 5
filters = 9
kernel = np.ones((kernel_size, kernel_size, 3, filters))
mask = np.ones(kernel.shape, dtype=np.float32)
mask[:,:,0,0]
mask[kernel_size // 2, kernel_size // 2 + 1:, :, :] = 0.
mask[kernel_size // 2 + 1:, :, :, :] = 0.
mask[:,:,0,0]
mask_type ='b'
in_channel = 3
if mask_type =='A':
    mask[kernel_size // 2, kernel_size // 2, :, :int(filters/3)] = 0.
    mask[kernel_size // 2, kernel_size // 2, 1:, int(filters/3):2*int(filters/3)] = 0.
    mask[kernel_size // 2, kernel_size // 2, 2:, 2*int(filters/3):] = 0.
else:
    mask[kernel_size // 2, kernel_size // 2, int(in_channel/3):, :int(filters/3)] = 0.
    mask[kernel_size // 2, kernel_size // 2, 2*int(in_channel/3):, int(filters/3):2*int(filters/3)] = 0.

mask[:,:,0,0]
class ResidualBlock(tf.keras.Model):
    """Residual blocks that compose pixelCNN

    Blocks of layers with 3 convolutional layers and one residual connection.
    Based on Figure 5 from [1] where h indicates number of filters.

    Refs:
    [1] - <NAME>., <NAME>., & <NAME>. (2016). Pixel
     recurrent neural networks. arXiv preprint arXiv:1601.06759.
    """

    def __init__(self, h):
        super(ResidualBlock, self).__init__(name='')

        self.conv2a = keras.layers.Conv2D(filters=h, kernel_size=1, strides=1)
        self.conv2b = MaskedConv2D(mask_type='B', filters=h, kernel_size=3, strides=1)
        self.conv2c = keras.layers.Conv2D(filters=2 * h, kernel_size=1, strides=1)

    def call(self, input_tensor):
        x = tf.nn.relu(input_tensor)
        x = self.conv2a(x)

        x = tf.nn.relu(x)
        x = self.conv2b(x)

        x = tf.nn.relu(x)
        x = self.conv2c(x)

        x += input_tensor
        return x
# Create PixelCNN model
inputs = keras.layers.Input(shape=(height, width, n_channel))
x = MaskedConv2D(mask_type='A', filters=128, kernel_size=7, strides=1)(inputs)

for i in range(15):
    x = ResidualBlock(h=64)(x)

x = keras.layers.Activation(activation='relu')(x)
x = keras.layers.Conv2D(filters=128, kernel_size=1, strides=1)(x)
x = keras.layers.Activation(activation='relu')(x)
x = keras.layers.Conv2D(filters=128, kernel_size=1, strides=1)(x)
x = keras.layers.Conv2D(filters=n_channel * q_levels, kernel_size=1, strides=1)(x)  # shape [N,H,W,DC]

pixelcnn = tf.keras.Model(inputs=inputs, outputs=x)
# Prepare optimizer and loss function
lr_decay = 0.9999
learning_rate = 5e-3
optimizer = tf.keras.optimizers.Adam(lr=learning_rate)

compute_loss = tf.keras.losses.CategoricalCrossentropy(from_logits=True)
@tf.function
def train_step(batch_x, batch_y):
    with tf.GradientTape() as ae_tape:
        logits = pixelcnn(batch_x, training=True)

        logits = tf.reshape(logits, [-1, height, width, q_levels, n_channel])  # shape [N,H,W,DC] -> [N,H,W,D,C]
        logits = tf.transpose(logits, perm=[0, 1, 2, 4, 3])  # shape [N,H,W,D,C] -> [N,H,W,C,D]

        loss = compute_loss(tf.one_hot(batch_y, q_levels), logits)

    gradients = ae_tape.gradient(loss, pixelcnn.trainable_variables)
    gradients, _ = tf.clip_by_global_norm(gradients, 1.0)
    optimizer.apply_gradients(zip(gradients, pixelcnn.trainable_variables))

    return loss
print("Num GPUs Available: ", len(tf.config.experimental.list_physical_devices('GPU')))


with tf.device("/job:localhost/replica:0/task:0/device:GPU:1"):
    # Training loop
    n_epochs = 50
    n_iter = int(np.ceil(x_train_quantised.shape[0] / batch_size))
    for epoch in range(n_epochs):
        start_epoch = time.time()
        for i_iter, (batch_x, batch_y) in enumerate(train_dataset):
            start = time.time()
            optimizer.lr = optimizer.lr * lr_decay
            loss = train_step(batch_x, batch_y)
            iter_time = time.time() - start
            if i_iter % 100 == 0:
                print('EPOCH {:3d}: ITER {:4d}/{:4d} TIME: {:.2f} LOSS: {:.4f}'.format(epoch,
                                                                                       i_iter, n_iter,
                                                                                       iter_time,
                                                                                       loss))
        epoch_time = time.time() - start_epoch
        print('EPOCH {:3d}: TIME: {:.2f} ETA: {:.2f}'.format(epoch,
                                                             epoch_time,
                                                             epoch_time * (n_epochs - epoch)))

A = np.zeros((1,10,10,3))
B = pixelcnn(A)
B = tf.reshape(B, [-1, height, width, q_levels, n_channel])
B = tf.transpose(B, perm=[0, 1, 2, 4, 3])
next_sample = tf.random.categorical(B[0,0,0,0,:][None,:], 1)
next_sample
# --------------------------------------------------------------------------------------------------------------
# Generating new images
samples = np.zeros((10, height,width,n_channel)) + + 0.5*np.random.rand(10,height,width, n_channel)

for i in range(height):
    for j in range(width):
        for k in range(n_channel):
            logits = pixelcnn(samples)
            logits = tf.reshape(logits, [-1, height, width, q_levels, n_channel])
            logits = tf.transpose(logits, perm=[0, 1, 2, 4, 3])
            next_sample = tf.random.categorical(logits[:, i, j, k, :], 1)
            samples[:, i, j, k] = (next_sample.numpy() / (q_levels - 1))[:, 0]
            

fig = plt.figure(figsize=(10, 10))
for i in range(9):
    ax = fig.add_subplot(3, 3, i+1)
    ax.imshow(samples[i, :, :, :])
    plt.xticks(np.array([]))
    plt.yticks(np.array([]))
plt.show()
# --------------------------------------------------------------------------------------------------------------
# Generating new images
occlude_start_row = 14
num_generated_images = 1
samples = np.copy(x_train_quantised[:10, :, :, :])
samples = samples / (q_levels - 1)
samples[:, occlude_start_row:, :, :] = 0

for i in range(occlude_start_row, height):
    for j in range(width):
        for k in range(n_channel):
            logits = pixelcnn(samples)
            logits = tf.reshape(logits, [-1, height, width, q_levels, n_channel])
            logits = tf.transpose(logits, perm=[0, 1, 2, 4, 3])
            next_sample = tf.random.categorical(logits[:, i, j, k, :], 1)
            samples[:, i, j, k] = (next_sample.numpy() / (q_levels - 1))[:, 0]
            

fig = plt.figure(figsize=(10, 10))
for i in range(9):
    ax = fig.add_subplot(3, 3, i+1)
    ax.imshow(samples[i, :, :, :])
    plt.xticks(np.array([]))
    plt.yticks(np.array([]))
plt.show()
# --------------------------------------------------------------------------------------------------------------
# Generating new images
occlude_start_row = 14
num_generated_images = 10
samples = np.copy(x_test_quantised[0:num_generated_images, :, :, :])
samples = samples / (q_levels - 1)
samples[:, occlude_start_row:, :, :] = 0

for i in range(occlude_start_row, height):
    for j in range(width):
        for k in range(n_channel):
            logits = pixelcnn(samples)
            logits = tf.reshape(logits, [-1, height, width, q_levels, n_channel])
            logits = tf.transpose(logits, perm=[0, 1, 2, 4, 3])
            next_sample = tf.random.categorical(logits[:, i, j, k, :], 1)
            samples[:, i, j, k] = (next_sample.numpy() / (q_levels - 1))[:, 0]
            

fig = plt.figure(figsize=(10, 10))
for i in range(9):
    ax = fig.add_subplot(3, 3, i+1)
    ax.imshow(samples[i, :, :, :])
    plt.xticks(np.array([]))
    plt.yticks(np.array([]))
plt.show()<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1983623877443212, 'fraction_numerical': 0.030771262546222925, 'mean_word_length': 2.281195840554593, 'pattern_counts': {'":': 0, '<': 8, '<?xml version=': 0, '>': 10, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 16, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '479595', 'license': 'MIT', 'n_tokens_mistral': 2840, 'n_tokens_neox': 2680, 'n_words': 550, 'path': 'WIP/2 - Modelling data with multiple channels/modelling_coloured_images.ipynb', 'repo_name': 'Mind-the-Pineapple/Autoregressive-models', 'size': '7596'}