Как реализовать t-SNE в тензорном потоке?

Я пытаюсь реализовать визуализацию t-SNE в тензорном потоке для задачи классификации изображений. То, что я в основном нашел в сети, было реализовано в Pytorch. См. здесь.

Вот мой общий код для обучения, который отлично работает, просто хочу добавить к нему визуализацию t-SNE:

import pandas as pd
import numpy as np
import tensorflow as tf
import cv2
from tensorflow import keras
from tensorflow.keras import layers, Input
from tensorflow.keras.layers import Dense, InputLayer, Flatten
from tensorflow.keras.models import Sequential, Model
from  matplotlib import pyplot as plt
import matplotlib.image as mpimg
from PIL import Image
from tensorflow.keras.preprocessing.image import ImageDataGenerator, array_to_img, img_to_array, load_img

.
.
.

base_model=tf.keras.applications.ResNet152(
    include_top=False, weights='imagenet', input_tensor=None,
    input_shape=None, pooling=None)

.
.
.

base_model.trainable = False


# Create new model on top.
inputs = tf.keras.Input(shape=(IMG_WIDTH, IMG_HEIGHT, 3))
x = base_model(inputs, training=False)



x=keras.layers.Flatten()(x)
x = keras.layers.Dense(64)(x)
x=layers.Activation('relu')(x)

x=keras.layers.Flatten()(x)
x = keras.layers.Dense(32)(x)
x=layers.Activation('relu')(x)

x = keras.layers.Dense(2)(x)
outputs=layers.Activation('softmax')(x)

model=keras.Model(inputs, outputs)



vaidation_datagen = ImageDataGenerator(rotation_range=90,
                                     zoom_range=0.2,
                                     horizontal_flip=True, 
                                     vertical_flip=True)


train_generator = train_datagen.flow_from_directory(
        train_path,  # this is the target directory
        target_size=target_size,  # all images will be resized to the target size
        color_mode='rgb',
        batch_size=batch_size,
        shuffle=True,
        class_mode='categorical',
        interpolation='nearest',
        seed=42)  # since we use binary_crossentropy loss, we need binary labels


validation_generator = vaidation_datagen.flow_from_directory(
        validation_path,  # this is the target directory
        target_size=target_size,  # all images will be resized to the target size
        color_mode='rgb',
        batch_size=batch_size,
        shuffle=True,
        class_mode='categorical',
        interpolation='nearest',
        seed=42)


model.compile(optimizer, loss , metrics)

model_checkpoint = tf.keras.callbacks.ModelCheckpoint((model_path+model_filename), monitor='val_loss',verbose=1, save_best_only=True)
model.summary()


history = model.fit(
     train_generator,
     steps_per_epoch = num_of_train_img_raw//batch_size,
     epochs = epochs, 
     validation_data = validation_generator, # relates to the validation data.
     validation_steps = num_of_val_img_raw//batch_size,
     callbacks=[model_checkpoint],
     use_multiprocessing = False)

Основываясь на предоставленной справочной ссылке, кажется, что мне нужно сначала сохранить функции, а оттуда применить t-SNE следующим образом (эта часть скопирована и вставлена ​​из здесь):

tsne = TSNE(n_components=2).fit_transform(features)


# scale and move the coordinates so they fit [0; 1] range
def scale_to_01_range(x):
    # compute the distribution range
    value_range = (np.max(x) - np.min(x))

    # move the distribution so that it starts from zero
    # by extracting the minimal value from all its values
    starts_from_zero = x - np.min(x)

    # make the distribution fit [0; 1] by dividing by its range
    return starts_from_zero / value_range

# extract x and y coordinates representing the positions of the images on T-SNE plot
tx = tsne[:, 0]
ty = tsne[:, 1]

tx = scale_to_01_range(tx)
ty = scale_to_01_range(ty)


# initialize a matplotlib plot
fig = plt.figure()
ax = fig.add_subplot(111)

# for every class, we'll add a scatter plot separately
for label in colors_per_class:
    # find the samples of the current class in the data
    indices = [i for i, l in enumerate(labels) if l == label]

    # extract the coordinates of the points of this class only
    current_tx = np.take(tx, indices)
    current_ty = np.take(ty, indices)

    # convert the class color to matplotlib format
    color = np.array(colors_per_class[label], dtype=np.float) / 255

    # add a scatter plot with the corresponding color and label
    ax.scatter(current_tx, current_ty, c=color, label=label)

# build a legend using the labels we set previously
ax.legend(loc='best')

# finally, show the plot
plt.show()

Буду признателен за помощь в соединении этих двух частей.