How to enable set device as GPU with Tensorflow?

Hi,

I am trying to figure out how I can set device to GPU with Tensorflow. With other frameworks I have been able to use GPU with the same instance. I have below code and error. It runs OOM with CPU instantly.

I know with Pytorch I can go as torch.cuda.set_device(), can you help me understand what’s the method in Tensorflow?

tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-uncased')
train_encodings = tokenizer(train_texts, truncation=True, padding=True)
train_dataset = tf.data.Dataset.from_tensor_slices((
    dict(train_encodings),
    train_labels
))

model = TFDistilBertForSequenceClassification.from_pretrained('distilbert-base-uncased', num_labels=3)

optimizer = tf.keras.optimizers.Adam(learning_rate=5e-5)
model.compile(optimizer=optimizer, loss= model.compute_loss, metrics=['accuracy'])

model.fit(train_dataset.shuffle(1000).batch(128), epochs=10, batch_size=128,
          validation_data=val_dataset.shuffle(1000).batch(128))

I get error as below:

2021-11-08 05:50:31.398734: W tensorflow/core/framework/op_kernel.cc:1745] OP_REQUIRES failed at transpose_op.cc:183 : RESOURCE_EXHAUSTED: OOM when allocating tensor with shape[128,512,12,64] and type float on /job:localhost/replica:0/task:0/device:CPU:0 by allocator cpu
1 Like

Normally, Keras automatically uses the GPU if it’s available.

You can check the available GPUs as follows:

import tensorflow as tf

print("Num GPUs Available: ", len(tf.config.list_physical_devices('GPU')))

Hey thanks for that. Seems that tensorflow actually doesn’t register my GPU when I run

tf.test.gpu_device_name()

I get error as below:

2021-11-08 22:08:10.636081: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:939] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2021-11-08 22:08:10.636549: W tensorflow/core/common_runtime/gpu/gpu_device.cc:1850] Cannot dlopen some GPU libraries. Please make sure the missing libraries mentioned above are installed properly if you would like to use GPU. Follow the guide at https://www.tensorflow.org/install/gpu for how to download and setup the required libraries for your platform.
Skipping registering GPU devices...
''

I can use the GPU with Pytorch definitely. When I change the code to Pytorch and place the model on GPU, and start training, I can see GPU usage increasing from nvidia-smi.

I’m using Tensorflow 2.7.0 as well, so it’s not any of the older versions where you need tensorflow-gpu as well.

I’m using cuda V11.0.221 and tried reinstalling but it seems that my account doesn’t have the permission to do that unfortunately.

Any ideas to tackle this would be appreciated highly :slight_smile:

Hi guys, I have trained my model-TFBertForSequenceClassification using GPU without any changes to the training code as given on the site and above by user. I can’t seem to figure out how to do inference for the same model on `GPU’?? The model.to_device(‘cuda’) doesn’t work in this case because of TF.

Any suggestions? Thanks