Skip to content

How to pretrain on GPU? #13

@ghost

Description

Hello!

I am trying to pretrain an adapter using the 4_pretrain_adapter.sh script.
I have a GeForce RTX 2080 SUPER installed (~8GB VRAM), with NVIDIA Driver Version: 440.33.01, CUDA Version: 10.2 and tensorflow-gpu 1.15.5.
I set CUDA_VISIBLE_DEVICES to 0 in the 4_pretrain_adapter.sh script since I only have a single GPU
Pretraining have been running for 12-16hrs now and is just about completing the warmup phase (~10000 steps).
I noticed that the pretraining only uses ~115MB of VRAM but several threads on CPU driving its usage up to ~100%.
I started perusing the code for GPU usage options/parameters, but, but so far only found a switch for TPU usage and a comment stipulating that if TPU is not available, then the Estimator (tf.contrib.tpu.TPUEstimator) will fall back on CPU or GPU.
I then looked at the official tensorflow documentation for TPUEstimator but no luck there either.

As I continue to look into this, I was wondering if you add some tips or advices about running the code locally on a single GPU.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions