How to train on only one GPU?

brain-research / self-attention-gan

Apache License 2.0

1k stars 174 forks source link

How to train on only one GPU? #12

Open carshadi opened 5 years ago

carshadi commented 5 years ago

Hello, Which parameters do I need to change to make this train and evaluate on one GPU? I am currently getting an OOM Resource Exhausted Error when i try to train on one GTX 1080. I tried setting num_towers=1 in train_imagenet.py but this did not help.

carshadi commented 5 years ago

Also tried reducing the batch size to 32,16,etc but the program got stuck before even starting to iterate

ArashHosseini commented 5 years ago

@carshadi, hi https://github.com/brain-research/self-attention-gan/blob/ad9612e60f6ba2b5ad3d3340ebae60f724636d75/train_imagenet.py#L124

did you try config.gpu_options.allow_growth=True ? or config.gpu_options.per_process_gpu_memory_fraction = 0.4