I have trained a neural network model which would generate the input data’s embeddings. Training is done easily.
Now, I wanted to generate the embeddings of the training dataset from the model. Currently, how I’ve implemented is that there is a dataloader for training dataset which is passed to the model to generate the embeddings. Once, model generates the embeddings, it is stored in a numpy array.
model.eval()
FX      = np.zeros((1,300)) # embedding size is 300
print('\nGenerating Embeddings..')

    with torch.no_grad():

        for bi, data in tqdm(enumerate(dataloader), total=len(dataloader)):

            ids            = data["ids"]
            mask           = data["mask"]
            token_type_ids = data["token_type_ids"]
            target         = data["targets"]        # This is 1-hot label row
            target         = target.squeeze(1)

            ids            = ids.to(config.device, dtype=torch.long)
            mask           = mask.to(config.device, dtype=torch.long)
            token_type_ids = token_type_ids.to(config.device, dtype=torch.long)

            output         = model(ids, mask, token_type_ids)

            FX             = np.append(FX, output.detach().cpu().numpy(), 0)

Problem I’m facing is as epoch moves ahead, it gets slower solely due to storing the model output in FX. I’m also detaching the model output from the computation graph and then converting it into numpy array but still GPU memory keeps rising and epoch gets slower after batch by batch.
I have tried commenting out FX             = np.append(FX, output.detach().cpu().numpy(), 0) and I found epoch doesn’t get slowed down. It worked perfectly.
How can I store the embeddings from the model such that the epoch doesn’t get slower?