<jupyter_start><jupyter_text>Load datasetsReminder: First, load, pre-process and store a dataset via the [RadialBeam](./radialbeamsampling.ipynb) routine.<jupyter_code>en_load_dataset = False

if en_load_dataset:
    train_dataset = tf.data.experimental.load('./data/{0}_train'.format(dataset_name)).batch(batch_size)
    val_dataset = tf.data.experimental.load('./data/{0}_val'.format(dataset_name)).batch(batch_size)
    test_dataset = tf.data.experimental.load('./data/{0}_test'.format(dataset_name)).batch(batch_size)
else:
    splits = [0.8, 0.1, 0.1]
    dataset = load_dataset(dataset_name)

    n_train = int(splits[0] * float(dataset.cardinality()))
    n_val = int(splits[1] * float(dataset.cardinality()))
    n_test = int(splits[2] * float(dataset.cardinality()))

    train_dataset = dataset.take(n_train)
    val_dataset = dataset.skip(n_train).take(n_val)
    test_dataset = dataset.skip(n_train).skip(n_val).take(n_train)

    img_size = int(train_dataset.element_spec['image'].shape[0])
    lines, angles = instantiate_radial_vectors(img_size + margin_padding, img_size + margin_padding,
                                               beam_set_size=n_beams,
                                               max_len=target_size)
    train_dataset = preprocess(train_dataset, lines, angles, target_size=img_size + margin_padding,
                               batch_size=batch_size, path='./training_dataset', continuous=continuous)
    val_dataset = preprocess(val_dataset, lines, angles, target_size=img_size + margin_padding,
                             batch_size=batch_size, path='./val_dataset', continuous=continuous)
    test_dataset = preprocess(test_dataset, lines, angles, target_size=img_size + margin_padding,
                              batch_size=batch_size, path='./test_dataset', continuous=continuous)

_, n_beams, _, n_pixels, n_channels = train_dataset.element_spec['beam'].shape<jupyter_output>2022-05-20 17:35:14.475394: I tensorflow/core/platform/cpu_feature_guard.cc:151] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 FMA
To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.
2022-05-20 17:35:15.138228: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1525] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 35742 MB memory:  -> device: 0, name: NVIDIA A40, pci bus id: 0000:43:00.0, compute capability: 8.6
2022-05-20 17:35:27.939613: W tensorflow/core/kernels/data/cache_dataset_ops.cc:768] The calling iterator did not fully read the dataset being cached. In order to avoid unexpected truncation of the dataset, the partially cached contents of the dataset  will be discarded. This can happen if you have an input pipeline similar to `dataset.cache().take(k).repeat()`. You should use `dataset.take(k).cache().rep[...]<jupyter_text>Load the BIC model<jupyter_code># input tensor (batch x (zero, theta) x beams x (2epsilon + 1) x D x C)
in_beams = tf.keras.layers.Input([2, n_beams, 3, n_pixels, n_channels])

bic = BIC(hidden=128, activation=tf.nn.leaky_relu, context=True,
          l2_regularization=0.0, edge_factor=0.5, gcn_layers=3, dropout=0.0,
          size_vector_field=n_beams, pixel_count_per_vector=n_pixels)

# multiple output for introspection; for training and inference: prior and unit_vec are essential
prior, unit_vec, beamencoding, ctx, similarity, \
beamencoding_zero, beamencoding_theta, angle_energy, rnn_encoding = bic(inputs=in_beams)

model = tf.keras.models.Model(inputs=in_beams, name='bic',
                              outputs=(prior, unit_vec, beamencoding, ctx, similarity, \
                                       beamencoding_zero, beamencoding_theta, angle_energy, rnn_encoding))

model = tf.keras.models.load_model(model_path)

model.summary()<jupyter_output>WARNING:tensorflow:No training configuration found in save file, so the model was *not* compiled. Compile it manually.<jupyter_text>Crop and canonicalize methods<jupyter_code>def remove_padding(image):
    cropped_image = tf.image.crop_to_bounding_box(image, margin_padding // 2, margin_padding // 2,
                                                  target_size * 2, target_size * 2)
    return cropped_image


def downstream_predict(downstream, sample, k=1, show=False):
    # predict non-rotated image with downstream model
    if k == 1:
        non_rot_pred = downstream(remove_padding(sample['image']))
    else:
        non_rot_pred = tf.nn.top_k(downstream(remove_padding(sample['image'])), k=k)

    # predict rotated and chopped off image
    if k == 1:
        crop_rot_pred = downstream(remove_padding(sample['rotated']))
    else:
        crop_rot_pred = tf.nn.top_k(downstream(remove_padding(sample['rotated'])), k=k)

    # call BIC
    pred_facts, pred_angle, conv_latents, gnn_latents, distance_matrix, \
    x1_emb, x2_emb, angle_energy, rnn_encoding = model(
        tf.tile(sample['beam_rot'][:, None, ...], [1, 2, 1, 1, 1, 1]))

    # project form complex vector to angle
    pred_angle = np.array([angle_between(pred_angle[b], tf.cast([1., 0.], tf.float32), gpu=True)
                           for b in range(tf.shape(pred_angle)[0])])

    # smoothly rotate the image back
    back_rot = tfa.image.rotate(sample['rotated'], 2 * math.pi - pred_angle,
                                interpolation='bilinear')

    # predict with downstream model
    if k == 1:
        if show:
            plt.imshow(remove_padding(back_rot[0]))
            plt.show()
        canonic_pred = downstream(remove_padding(back_rot))
    else:
        canonic_pred = tf.nn.top_k(downstream(remove_padding(back_rot)), k=k)

    return non_rot_pred, crop_rot_pred, canonic_pred
cce = tf.keras.losses.CategoricalCrossentropy(from_logits=False, label_smoothing=0, axis=-1)
train_df = pd.DataFrame(columns=['iteration', 'non_rot_loss', 'crop_rot_loss', 'canonic_loss'])
valid_df = pd.DataFrame(columns=['iteration', 'non_rot_loss', 'crop_rot_loss', 'canonic_loss'])
n_runs = 5
epochs = 8

for _ in range(n_runs):
    downstream = VGG11((target_size * 2, target_size * 2, 3), 100)
    optimizer = tf.keras.optimizers.Adam(learning_rate=0.001, beta_1=0.9, beta_2=0.999, epsilon=1e-07)
    for e in tqdm(range(epochs)):

        # validation, before training to log also the init behaviour of the model
        for i, sample in enumerate(val_dataset):
            non_rot_pred, crop_rot_pred, canonic_pred = downstream_predict(downstream, sample)
            valid_df = valid_df.append(pd.DataFrame({
                'iteration': [e * int(val_dataset.cardinality().numpy()) + i],
                'non_rot_loss': [float(cce(tf.one_hot(sample['label'], 100), non_rot_pred).numpy())],
                'crop_rot_loss': [float(cce(tf.one_hot(sample['label'], 100), crop_rot_pred).numpy())],
                'canonic_loss': [float(cce(tf.one_hot(sample['label'], 100), canonic_pred).numpy())]
            }), ignore_index=True)
            pass

        # training
        for i, sample in enumerate(train_dataset):
            with tf.GradientTape() as tape:
                non_rot_pred, crop_rot_pred, canonic_pred = downstream_predict(downstream, sample,
                                                                               show=False)#True if i == 0 else False)
                loss = cce(tf.one_hot(sample['label'], 100), non_rot_pred)
                train_df = train_df.append(pd.DataFrame({
                    'iteration': [e * int(train_dataset.cardinality().numpy()) + i],
                    'non_rot_loss': [float(cce(tf.one_hot(sample['label'], 100), non_rot_pred).numpy())],
                    'crop_rot_loss': [float(cce(tf.one_hot(sample['label'], 100), crop_rot_pred).numpy())],
                    'canonic_loss': [float(cce(tf.one_hot(sample['label'], 100), canonic_pred).numpy())]
                }), ignore_index=True)
            grads_downstream = tape.gradient(loss, downstream.trainable_variables)
            optimizer.apply_gradients(zip(grads_downstream, downstream.trainable_variables))
colors = ['#629FCA', '#FDA556', '#6BBC6B', '#E26768', '#B292CE']

ax = sns.lineplot(data=valid_df, x='iteration', y='non_rot_loss', label='non_rot_loss', ci='sd',
             palette=colors[0], alpha=0.7)
ax = sns.lineplot(data=valid_df, x='iteration', y='crop_rot_loss', label='crop_rot_loss', ci='sd',
             palette=colors[1], alpha=0.7)
ax = sns.lineplot(data=valid_df, x='iteration', y='canonic_loss', label='canonic_loss', ci='sd',
             palette=colors[2], alpha=0.7)

ax.set_ylabel('Cross Entropy')
ax.legend().remove()
plt.show()<jupyter_output><empty_output>