Визуализация встраивания узлов последнего слоя модели в факел геометрического

Я делаю свой первый проект графовой сверточной нейронной сети с torch_geometric. Я хочу визуализировать вложения узлов последнего слоя моей модели и не знаю, как мне это получить.

Я обучил свою модель набору данных CiteSeer. Вы можете получить полный набор данных так же просто, как это:

from torch_geometric.datasets import Planetoid
from torch_geometric.transforms import NormalizeFeatures
dataset = Planetoid(root = "data/Planetoid", name='CiteSeer', transform=NormalizeFeatures())

Моя модель представляет собой простую двухслойную модель:

class GraphClassifier(torch.nn.Module):
    def __init__(self, dataset, hidden_dim):
        super(GraphClassifier, self).__init__()
        self.conv1 = GCNConv(dataset.num_features, hidden_dim)
        self.conv2 = GCNConv(hidden_dim, dataset.num_classes)

    def forward(self, data):
        x, edge_index = data.x, data.edge_index
        x = F.relu(self.conv1(x, edge_index))
        x = F.relu(self.conv2(x, edge_index))
        return F.log_softmax(x, dim=1) 

Если вы распечатаете мою модель, вы получите это:

model = GraphClassifier(dataset, 64)
print(model)

>>>
GraphClassifier(
  (conv1): GCNConv(3703, 64)
  (conv2): GCNConv(64, 6)
)

Моя модель успешно обучена. Я только хочу визуализировать его вложения узлов последнего уровня. Чтобы визуализировать, что у меня есть эта функция для использования:

%matplotlib inline
import matplotlib.pyplot as plt
from sklearn.manifold import TSNE
import torch
# emb: (nNodes, hidden_dim)
# node_type: (nNodes,). Entries are torch.int64 ranged from 0 to num_class - 1

def visualize(emb: torch.tensor, node_type: torch.tensor):
  z = TSNE(n_components=2).fit_transform(emb.detach().cpu().numpy())
  plt.figure(figsize=(10,10))
  plt.scatter(z[:, 0], z[:, 1], s=70, c=node_type, cmap = "Set2")
  plt.show()

Я не знаю, как мне извлечь emb и node_type из моей модели, чтобы передать их функции visualize. emb — последний слой вложений узлов модели. Как я могу получить их от моей модели?