Jądro Multi-GPU Tensorflow Java

Question

Mam serwer z wieloma procesorami graficznymi i chcę w pełni z nich korzystać podczas wnioskowania modelu wewnątrz aplikacji java. Domyślnie tensorflow zajmuje wszystkie dostępne GPU, ale używa tylko pierwszego.

mogę myśleć o trzech opcji rozwiązania tego problemu:

1. Ogranicz widoczność urządzenia na poziomie procesu, a mianowicie za pomocą CUDA_VISIBLE_DEVICES zmienną środowiskową.

   Wymagałoby to uruchomienia kilku wystąpień aplikacji Java i rozpowszechniania ruchu między nimi. Nie ten kuszący pomysł.

2. uruchomienie kilku sesji wewnątrz jednej aplikacji i spróbować przypisać jedno urządzenie do każdego z nich poprzez 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 = argmin(counters); 
        counters[i] += 1; 
        return i; 
       } 
   
       private synchronized void releasePredictorIndex(int i) { 
        counters[i] -= 1; 
       } 
   } 
   
   
   public class Predictor { 
   
       private Session session; 
   
       public Predictor(String modelPath, int deviceIdx, int numDevices, int numThreadsPerDevice) { 
   
        GPUOptions gpuOptions = GPUOptions.newBuilder() 
          .setVisibleDeviceList("" + deviceIdx) 
          .setAllowGrowth(true) 
          .build(); 
   
        ConfigProto config = ConfigProto.newBuilder() 
          .setGpuOptions(gpuOptions) 
          .setInterOpParallelismThreads(numDevices * numThreadsPerDevice) 
          .build(); 
   
        byte[] graphDef = Files.readAllBytes(Paths.get(modelPath)); 
        Graph graph = new Graph(); 
        graph.importGraphDef(graphDef); 
   
        this.session = new Session(graph, config.toByteArray()); 
       } 
   
       public Prediction predict(Data data) { 
        // ... 
       } 
   } 
   

   Takie podejście wydaje się działać dobrze na pierwszy rzut oka. Jednak sesje sporadycznie ignorują opcję setVisibleDeviceList i wszystkie idą na pierwsze urządzenie powodujące awarię z powodu braku pamięci.

3. Zbuduj model w stylu wielopiętrowym w pythonie, używając specyfikacji tf.device(). Po stronie Java podaj różne wieże wewnątrz współdzielonej sesji.

   Czuje się uciążliwy i idiomatycznie zły dla mnie.

UPDATE: jako @ash zaproponowano tam jeszcze innej opcji:

1. przypisać odpowiednie urządzenie dla każdej operacji istniejącego wykresu poprzez zmianę określenia (graphDef).

   Aby to zrobić, można dostosować kod Metoda 2:

   public class Predictor { 
   
       private Session session; 
   
       public Predictor(String modelPath, int deviceIdx, int numDevices, int numThreadsPerDevice) { 
   
        byte[] graphDef = Files.readAllBytes(Paths.get(modelPath)); 
        graphDef = setGraphDefDevice(graphDef, deviceIdx) 
   
        Graph graph = new Graph(); 
        graph.importGraphDef(graphDef); 
   
        ConfigProto config = ConfigProto.newBuilder() 
          .setAllowSoftPlacement(true) 
          .build(); 
   
        this.session = new Session(graph, config.toByteArray()); 
       } 
   
       private static byte[] setGraphDefDevice(byte[] graphDef, int deviceIdx) throws InvalidProtocolBufferException { 
        String deviceString = String.format("/gpu:%d", deviceIdx); 
   
        GraphDef.Builder builder = GraphDef.parseFrom(graphDef).toBuilder(); 
        for (int i = 0; i < builder.getNodeCount(); i++) { 
         builder.getNodeBuilder(i).setDevice(deviceString); 
        } 
        return builder.build().toByteArray(); 
       } 
   
       public Prediction predict(Data data) { 
        // ... 
       } 
   } 
   

   Podobnie jak w innych wymienionych podejść, ten nie uwolnij mnie od ręcznie dystrybucji danych między urządzeniami. Ale przynajmniej działa stabilnie i jest porównywalnie łatwa do wdrożenia. Ogólnie rzecz biorąc, wygląda to jak (prawie) normalna technika.

Czy istnieje elegancki sposób na wykonanie tak podstawowej rzeczy za pomocą interfejsu API tensorflow java? Wszelkie pomysły będą mile widziane.

Answer 1

Krótko mówiąc: istnieje obejście, w którym kończy się jedna sesja na układ GPU.

Dane:

Główny przepływ jest to, że czas pracy TensorFlow uwzględnia urządzeń przeznaczonych dla operacji na wykresie. Jeśli dla operacji nie określono żadnego urządzenia, to "umieszcza" je na podstawie heurystyki. Heurystyki te powodują obecnie "miejsce pracy na GPU: 0, jeśli GPU są dostępne i jest jądro GPU dla operacji" (Placer::Run w razie zainteresowania).

To, o co pytasz, to według mnie rozsądne żądanie funkcji dla TensorFlow - możliwość traktowania urządzeń w serializowanym wykresie jako "wirtualnych", które mają być odwzorowane w zestawie urządzeń "phyiscal" w czasie wykonywania, lub alternatywnie ustawienie "urządzenie domyślne". Ta funkcja obecnie nie istnieje. Dodanie takiej opcji do ConfigProto jest czymś, do czego możesz chcieć złożyć żądanie funkcji.

Mogę zaproponować tymczasowe obejście tego problemu. Po pierwsze, komentarz do twoich proponowanych rozwiązań.

1. Twój pierwszy pomysł na pewno zadziała, ale jak zauważyłeś, jest kłopotliwy.

2. Ustawienie przy użyciu visible_device_list w ConfigProto nie całkiem działa, ponieważ jest to ustawienie na proces i jest ignorowane po utworzeniu pierwszej sesji w procesie. Z pewnością nie jest to udokumentowane, a powinno być (i trochę niefortunne, że pojawia się w konfiguracji na sesję). Wyjaśnia to jednak, dlaczego Twoja sugestia nie działa i dlaczego nadal używasz jednego procesora graficznego.

3. To może działać.

Inną możliwością jest, aby zakończyć się przy różnych wykresów (z operacji wyraźnie umieszczone na różnych GPU), w wyniku czego w trakcie jednej sesji za graficznego. Coś takiego może być używany do edycji wykresu i jednoznacznie przypisać urządzenie do każdej operacji:

public static byte[] modifyGraphDef(byte[] graphDef, String device) throws Exception { 
    GraphDef.Builder builder = GraphDef.parseFrom(graphDef).toBuilder(); 
    for (int i = 0; i < builder.getNodeCount(); ++i) { 
    builder.getNodeBuilder(i).setDevice(device); 
    } 
    return builder.build().toByteArray(); 
} 

Po czym można utworzyć Graph i Session za GPU przy użyciu coś jak:

final int NUM_GPUS = 8; 
// setAllowSoftPlacement: Just in case our device modifications were too aggressive 
// (e.g., setting a GPU device on an operation that only has CPU kernels) 
// setLogDevicePlacment: So we can see what happens. 
byte[] config = 
    ConfigProto.newBuilder() 
     .setLogDevicePlacement(true) 
     .setAllowSoftPlacement(true) 
     .build() 
     .toByteArray(); 
Graph graphs[] = new Graph[NUM_GPUS]; 
Session sessions[] = new Session[NUM_GPUS]; 
for (int i = 0; i < NUM_GPUS; ++i) { 
    graphs[i] = new Graph(); 
    graphs[i].importGraphDef(modifyGraphDef(graphDef, String.format("/gpu:%d", i))); 
    sessions[i] = new Session(graphs[i], config);  
} 

Następnie użyj sessions[i], aby wykonać wykres na GPU #i.

Nadzieję, że pomaga.