Tôi có một máy chủ có nhiều GPU và muốn tận dụng chúng trong khi suy luận mô hình bên trong ứng dụng java. Theo mặc định, lưu lượng sẽ thu thập tất cả các GPU hiện có, nhưng chỉ sử dụng GPU đầu tiên.Suy luận Multi-GPU của dòng chảy Tensor
tôi có thể nghĩ ra ba lựa chọn để vượt qua vấn đề này:
Giới hạn tầm nhìn thiết bị trên mức độ xử lý, cụ thể là sử dụng
CUDA_VISIBLE_DEVICESbiến môi trường.Điều đó sẽ yêu cầu tôi chạy một số phiên bản của ứng dụng java và phân phối lưu lượng truy cập trong số đó. Không phải là ý tưởng hấp dẫn đó.
Launch nhiều phiên bên trong một ứng dụng duy nhất và cố gắng gán một thiết bị cho mỗi trong số họ qua
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) { // ... } }Cách tiếp cận này dường như làm việc tốt trong nháy mắt. Tuy nhiên, các phiên thỉnh thoảng bỏ qua tùy chọn
setVisibleDeviceListvà tất cả đều chuyển sang thiết bị đầu tiên gây ra lỗi ngoài bộ nhớ.Tạo mô hình theo kiểu đa tháp ở trăn sử dụng đặc điểm
tf.device(). Về phía java, hãy cung cấp các tòa tháp khác nhau củaPredictortrong một phiên được chia sẻ.Cảm thấy rườm rà và vô lý đối với tôi.
UPDATE: Như @ash đề xuất, có thêm một lựa chọn:
Gán một thiết bị phù hợp với từng hoạt động của đồ thị hiện bằng cách sửa đổi định nghĩa của nó (
graphDef).Để làm cho nó thực hiện, người ta có thể thích ứng với các mã từ Cách 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) { // ... } }Cũng giống như cách tiếp cận nêu trên khác, một trong những điều này không giải phóng tôi khỏi tay phân phối dữ liệu giữa các thiết bị. Nhưng ít nhất nó hoạt động ổn định và tương đối dễ thực hiện. Nhìn chung, điều này trông giống như một kỹ thuật (gần như) bình thường.
Có cách nào thanh lịch để thực hiện điều cơ bản như vậy với API java lưu lượng không? Có những câu chuyện mới trên trang chủ.