Giters

idealo / image-quality-assessment

Convolutional Neural Networks to predict the aesthetic and technical quality of images.

Home Page:https://idealo.github.io/image-quality-assessment/

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

Use specific GPU

x1ng4me opened this issue · comments

commented

I am trying to use this project as a part of my whole project. What I want is to run this part on a specific gpu. I found the part which used most gpu storage is the BaseCnn part in the model_builder.py.
self.base_model = BaseCnn(input_shape=(224, 224, 3), weights=self.weights, include_top=False, pooling='avg')
I am wondering if there is an api to let me specific the gpu number I want to use for this line.

ps: os.environ['CUDA_VISIBLE_DEVICES'] = "1" does not work for me, because the predict.py file is not the main script of my project.