Retrieve Embeddings of Samples in PatchTST

[MMayr96]

Given X of shape [<n_samples>, <n_features>, <n_lookback>] and a trained TSForecaster learner based on PatchTST, is there a way to retrieve the embeddings for each of the samples?

I found a related discussion that was closed Dec. 2022 for TST (see #560 ), however I have troubles finding the right shape when triggering the forward() of the PatchTST model. As PatchTST operates on patches, the forward implementation operates on the following shape:
x = x.permute(0,1,3,2) # x: [bs x nvars x patch_num x patch_len]

Does anyone have some suggestions on how to tackle this problem and retrieve the embeddings?

PS: (I also modified the suggested method in the referenced discussion to fit PatchTST. I am not quite sure if this is correct but it seems to hang into the right place for embedding retrieval):
net = nn.Sequential(learn.model.model.backbone,list(list(learn.model.model.head.children())[0].children())[:-1][0])

[oguiza]

Hi @MMayr96,
If by embeddings you mean the output of a given layer, I'd recommend you use fastai's Hook (see this to learn more about hooks).
You just need to decide which output you want to pull. If you want to pull the output of the backbone before it's passed to the model´s head, you could use something like this:

1. Instantiate a learner as you would normally do:

from tsai.basics import *
from fastai.callback.hook import Hook

fcst_history = 60
fcst_horizon = 3

df = get_forecasting_time_series("Sunspots")
X, y = prepare_forecasting_data(df, fcst_history=fcst_history, fcst_horizon=fcst_horizon, dtype=np.float32)
splits = get_forecasting_splits(df, fcst_history=fcst_history, fcst_horizon=fcst_horizon, test_size=235, show_plot=False)
batch_tfms = TSStandardize()
fcst = TSForecaster(X, y, splits=splits, path='models', batch_tfms=batch_tfms, bs=64, arch="PatchTST", metrics=mae, cbs=ShowGraph())

2. To grab the ambeddings from a single batch you could use:

hook = Hook(fcst.model.model.backbone, lambda m,i,o: o, is_forward=True, detach=True, cpu=True)
xb,_ = fcst.dls.train.one_batch()
_ = fcst.model(xb)
embeddings = hook.stored

3. To grab all embeddings you can use a loop:

hook = Hook(fcst.model.model.backbone, lambda m,i,o: o, is_forward=True, detach=True, cpu=True)
embeddings = []
for xb,_ in fcst.dls.train:
    _ = fcst.model(xb)
    embeddings.append(hook.stored)
embeddings = torch.cat(embeddings)

Is that what you were looking for?

[MMayr96]

Thank you very much for the quick and clear response, that is perfect!!