Tensorflow GPU memory allocation issue

51 Views Asked by At

I found something interesting: With TensorFlow, the code does not run when the memory is not sufficient to allocate on Windows and WSL2; it runs on MacBook even though it is very slow. Is anyone aware of this?

I was trying to run a text classification task and testing different batch sizes. Now I'm running it on Google Colab but I wish I could run it on my local machine with 1 GPU (RTX3050 or GTX1060Ti).

I ran the same code on Windows, WSL2, and Macbook Pro 13 inch with M1 processor. Windows and WSL2 are on the same machine.

Edited; I added a code for your reference.

from transformers import create_optimizer
import tensorflow as tf

batch_size = 16
num_epochs = 3
batches_per_epoch = len(tokenized_imdb["review"]) // batch_size
total_train_steps = int(batches_per_epoch * num_epochs)
optimizer, schedule = create_optimizer(init_lr=2e-5, num_warmup_steps=0, num_train_steps=total_train_steps)

from transformers import TFAutoModelForSequenceClassification

model = TFAutoModelForSequenceClassification.from_pretrained(
    "distilbert-base-uncased")

tf_train_set = model.prepare_tf_dataset(
    tokenized_imdb,
    shuffle=True,
    batch_size=32,
    collate_fn=data_collator,
)

tf_validation_set = model.prepare_tf_dataset(
    tokenized_imdb,
    shuffle=False,
    batch_size=16,
    collate_fn=data_collator,
)

tf_test_set = model.prepare_tf_dataset(
    tokenized_imdb,
    shuffle=False,
    batch_size=16,
    collate_fn=data_collator,
)

model.fit(x=tf_train_set, validation_data=tf_validation_set, epochs=3, callbacks=callbacks)
0

There are 0 best solutions below