小编Ale*_*nko的帖子

我有一个具有多个GPU的服务器,并希望在Java应用程序内的模型推理期间充分利用它们.默认情况下,tensorflow会占用所有可用的GPU,但仅使用第一个GPU.

我可以想出三个选项来克服这个问题:

1. 限制设备在流程级别上的可见性,即使用CUDA_VISIBLE_DEVICES环境变量.

   这将要求我运行java应用程序的几个实例并在它们之间分配流量.不是那个诱人的想法.

2. 在单个应用程序中启动多个会话,并尝试通过ConfigProto以下方式为每个会话分配一个设备:

   public class DistributedPredictor {
   
       private Predictor[] nested;
       private int[] counters;
   
       // ...
   
       public DistributedPredictor(String modelPath, int numDevices, int numThreadsPerDevice) {
           nested = new Predictor[numDevices];
           counters = new int[numDevices];
   
           for (int i = 0; i < nested.length; i++) {
               nested[i] = new Predictor(modelPath, i, numDevices, numThreadsPerDevice);
           }
       }
   
       public Prediction predict(Data data) {
           int i = acquirePredictorIndex();
           Prediction result = nested[i].predict(data);
           releasePredictorIndex(i);
           return result;
       }
   
       private synchronized int acquirePredictorIndex() {
           int i …

   Run Code Online (Sandbox Code Playgroud)