Function AvailableGPUMemory

tensorflow/core/grappler/devices.cc:69–82

}

int64 AvailableGPUMemory(int gpu_id) {
#if GOOGLE_CUDA
  // Look up the device, to see its attributes.
  se::Platform* gpu_platform = GPUMachineManager();
  CHECK_LT(gpu_id, gpu_platform->VisibleDeviceCount());
  se::StreamExecutor* se = gpu_platform->ExecutorForDevice(gpu_id).ValueOrDie();
  int64 total_memory, available_memory;
  CHECK(se->DeviceMemoryUsage(&available_memory, &total_memory));

  return available_memory;
#else
  return 0;
#endif
}

int GetNumAvailableLogicalCPUCores() { return port::NumSchedulableCPUs(); }