Как обрабатывать данные партиями в Tensorflow 2.0
У меня есть данные в формате csv. Я получаю данные из csv с помощью read_csv, а затем предварительно обрабатываю текст для DistilBERT - токенизация, а затем заполнение:
train_csv = pd.read_csv(train_csv)
train_df = train_csv[["text", "label"]]
BATCH_SIZE = 256
tokenized = train_df["text"].apply((lambda x:tokenizer.encode(x, add_special_tokens=True)))
padded = tf.keras.preprocessing.sequence.pad_sequences(tokenized.values, padding="post", maxlen=512)
attention_mask = np.where(padded != 0, 1, 0)
input_ids = torch.tensor(padded)
attention_mask = torch.tensor(attention_mask)
with torch.no_grad():
last_hidden_states = model(input_ids, attention_mask=attention_mask)
при выполнении последнего шага у меня возникают проблемы с оперативной памятью. может кто-нибудь, пожалуйста, поделитесь, как обрабатывать это партиями и получать последнее скрытое состояние?
Спасибо.