1. 首页 > 生活百科

不再让CPU和总线拖后腿:Exafunction让GPU跑的更快 不再让cpu和总线连接

作者:admin 更新时间:2025-09-06
摘要:对于并行运算,GPU 的应用效率是最高的。在云服务中使用 GPU 是获得低延迟..." />

对于并行运算,GPU 的应用效率是顶尖的。

在云服务中运用 GPU 是获取低延迟深度进修推理服务最经济的方法。运用 GPU 的主要瓶颈其中一个是通过 PCIe 总线在 CPU 和 GPU 内存之间复制数据的速度。对于许多打算用于高分辨率图像和视频处理的深度进修模型来说,简单地复制输入会大大增加体系的整体延迟,特别是当非推理任务,如解压缩和预处理也可以在 GPU 上执行时。

在这篇博文中,研究者们将展示怎样在 TensorFlow 中直接通过 GPU 内存传递模型输入和输出以进行模型推理,完全绕过 PCIe 总线和 CPU 内存。

由于大多数 GPU 代码是用 CUDA 编写的,这篇文章小编将将运用 TensorFlow 的 C++ 接口来示范这种技术。这样有利于对接其他库的接口,如用于 GPU 加速的图像预处理的 OpenCV 和用于硬件加速的视频解码的 NVIDIA NVDEC。

初始配置

TensorFlow 的 C++ 接口中,tensorflow::LoadSavedModel 被用来加载模型包:

tensorflow::SavedModelBundle bundle;
TF_RETURN_IF_ERROR(bundle, tensorflow::LoadSavedModel(
    session_options, run_options, saved_model_dir, tags, &bundle
));

接着可以运用 tensorflow::Session 来运行模型包。默认情况下,这将运用 CPU。

tensorflow::Session* session = bundle.GetSession();

// Create a tensor in CPU memory
tensorflow::Tensor tensor(tensorflow::DT_FLOAT, {1, 2, 3});
// Pairs of feed name and tensor to pass into the model
std::vector> inputs{"input", std::move(tensor)};

// The outputs will be written here. These will also be on the CPU.
std::vector outputs;
// Tensor names to fetch for the output.
std::vector fetch_names;

// Run the model!
session->Run(inputs, fetch_names, {}, &outputs);

运用 GPU

运用 GPU 就相对麻烦了。首先,用户必须从会话中创建壹个 tensorflow::ofableOptions 的实例,以指定哪些张量被传入和传出 GPU 内存而不是 CPU 内存。除了这些之后,有必要指定内存将从何者 GPU 中输入和获取。在这个例子中,为了简单起见,这篇文章小编将将把全部的输入和输出的张量(Tensor)放在第壹个 GPU 上。

tensorflow::使命召唤17ableOptions callable_options;
std::string gpu_device_name = FirstGpuDeviceName(session);

// Names of input tensors.
std::vector feed_names;
*callable_options.mutable_feed() = {feed_names.begin(), feed_names.end()};
// Names of output tensors.
std::vector fetch_names;
*callable_options.mutable_fetch() = {fetch_names.begin(), fetch_names.end()};

auto& feed_devices = *callable_options.mutable_feed_devices();
for (auto& input_name : feed_names) {
  feed_devices[input_name] = gpu_device_name;
}
auto& fetch_devices = *callable_options.mutable_fetch_devices();
for (auto& output_name : fetch_names) {
  fetch_devices[output_name] = gpu_device_name;
}

运用下面的函数可以获取 GPU 设备的名称:

std::string FirstGpuDeviceName(tensorflow::Session* session) {
  // Gets device name for the first GPU in the session.
  std::vector devices;
  auto status = session->ListDevices(&devices);
  assert(status.ok());
  for (const tensorflow::DeviceAttributes& d : devices) {
    if (d.device_type() == "GPU" || d.device_type() == "gpu") {
      return d.name();
    }
  }
  CHECK(false) << "GPU not found";
}

现在,用户可以创建壹个 tensorflow::Session::ofableHandle 的实例,这个类封装了怎样在 GPU 上运行带有输入和输出的 TensorFlow 图的方式。创建和销毁可调用对象的代价相对大,因此最好只在模型初始化时创建和销毁可调用对象。另外,可调用的对象应该在会话本身被销毁之前被销毁。

TF_RETURN_IF_ERROR(session->MakeWarable(callable_options, &callable));
// Before the session is destroyed:
// session->Releaseofable(callable);

接下来就可以创建一些输入张量了。在这个例子中,这篇文章小编将将只运用 TensorFlow 内置的 GPU 分配器,但其实也是可以通过 tensorflow::TensorBuffer 接口将外部张量传入外部 GPU 缓冲区。

// Get TensorFlow's GPU allocator for device 0
// This needs to match the device placement used when loading the SavedModel
// and creating the session.
tensorflow::TfDeviceId gpu_device_id(0);
tensorflow::Allocator* gpu_allocator =
    tensorflow::GPUProcessState::singleton()->GetGPUAllocator(gpu_device_id);

// Synchronize to ensure memory can be safely allocated & overwritten
cudaDeviceSynchronize();

// The input tensors are now allocated in GPU memory using TensorFlow's
// allocator. They must be in the same order as the feed names.
std::vector inputs;
for (int i = 0; i < 10; i++) {
  tensorflow::Tensor tensor(gpu_allocator, tensorflow::DT_FLOAT, {1, 2, 3});
  // Fill the input here
  inputs.push_back(std::move(tensor));
}

// Synchronize to ensure the inputs are valid
cudaDeviceSynchronize();

最后就可以运行模型了。现在,TensorFlow 既可以直接运用来自 GPU 的输入,也可以将输出放在同壹个 GPU 上

// The outputs will also be placed on the GPU thanks to the fetch_devices
// setting above.
std::vector outputs;

TF_RETURN_IF_ERROR(session->RunColdable(callable, inputs, &outputs, nullptr));

运用 CUDA stream

虽然 TensorFlow 内部运用 CUDA stream,但上述样例中全部的 CUDA 操作仍然是同步的。运行 cudaDeviceSynchronize 必须要在分配内存之前,以确保不会破坏先前分配好的 TensorFlow 内存。还必须在写入输入后进行同步操作,以确保 TensorFlow 能获取到有效的输入。TensorFlow 本身也会在模型执行结束时和 GPU 进行同步,以确保输出的张量是有效的。

显然,大众希望 GPU 能尽也许长时刻地异步运行以减少 CPU 造成的阻塞。幸运的是,用户可以访问内部的 TensorFlow CUDA stream。TensorFlow CUDA stream 的输入必须和 TensorFlow 的流同步,而输出的运用对象必须在访问内存之前和 TensorFlow 的流同步。运用 TensorFlow CUDA stream,大家可以完全关掉和 CPU 的同步。

具体来说,首先,在 OpsableOptions 上配置壹个额外的选项,以便在模型执行结束时禁用 TensorFlow 的内部同步。

callable_options.set_fetch_skip_sync(true);

可以运用下面的辅助函数访问内部流,需要注意的是参数包括设备名称。

cudaStream_t stream = GetTfGpuStream(session, gpu_device_name);

// Returns the tensorflow::BaseGPUDevice for a given device name
tensorflow::BaseGPUDevice* GetTfGpuDevice(tensorflow::Session* session,
                                          const std::string& gpu_device_name) {
  const tensorflow::DeviceMgr* device_mgr;
  auto status = session->LocalDeviceManager(&device_mgr);
  CHECK(status.ok()) << status;
  tensorflow::Device* device;
  status = device_mgr->LookupDevice(gpu_device_name, &device);
  CHECK(status.ok()) << status;
  auto* gpu_device = dynamic_cast(device);
  return CHECK_NOTNULL(gpu_device);
}

// Returns the compute stream for a given TensorFlow GPU device.
cudaStream_t GetTfGpuStream(tensorflow::Session* session,
                            const std::string& gpu_device_name) {
  auto* device = GetTfGpuDevice(session, gpu_device_name);
  const tensorflow::DeviceBase::GpuDeviceInfo* device_info =
      device->tensorflow_gpu_device_info();
  CHECK_NOTNULL(device_info);
  CUstream tf_stream =
      stream_executor::gpu::AsGpuStreamValue(device_info->stream);
  return tf_stream;
}

创建模型的输入,并如下面的代码所示运行。注意这里没有调用 cudaDeviceSynchronize!

cudaStream_t stream = GetTfGpuStream(session, gpu_device_name);

std::vector inputs;
for (int i = 0; i < 10; i++) {
  tensorflow::Tensor tensor(gpu_allocator, tensorflow::DT_FLOAT, {1, 2, 3});
  // Fill input buffers here, using the stream in kernel calls
  // or calls to cudaMemcpyAsync. Alternatively, synchronize with another
  // GPU stream using events.
  inputs.push_back(std::move(tensor));
}

std::vector outputs;
TF_RETURN_IF_ERROR(session->RunDutyable(callable, inputs, &outputs, nullptr));

cudaStreamSynchronize(stream);

请注意,如果 TensorFlow 内部需要将内存从 GPU 复制到 CPU,那么在运行模型时仍然也许发生 CPU 和 GPU 同步。然而,在给模型传递输入和输出时不再固有的需要任何和 CPU 的同步。

结论

作者旨在通过这篇文章示范怎样只通过 GPU 将输入和输出传递给 TensorFlow,这样一来可以绕过 PCIe 总线,减少开销和有限的 CPU 内存带宽。在 Exafunction,研究者们将在他们的模型服务化解方法——ExaDeploy——中运用这样的技术,以最大限度地进步 GPU 的利用率,即使是那些具有特别大的输入和输出的模型。

参考内容:

https://exafunction.com/blog/tensorflow-gpu-inputs-outputs