torch.nn.functional.torch.nn.parallel.data_parallel

torch.nn.parallel.data_parallel(module, inputs, device_ids=None, output_device=None, dim=0, module_kwargs=None)

在給定的 device_ids 的 GPU 上並行評估 module(input)。

這是 DataParallel 模組的函式式版本。

引數

• module (Module) – 要並行評估的模組

• inputs (Tensor) – 模組的輸入

• device_ids (list of int or torch.device) – 複製模組的 GPU ID

• output_device (list of int or torch.device) – 輸出的 GPU 位置。使用 -1 表示 CPU。(預設值: device_ids[0])

返回

一個位於 output_device 的 Tensor，包含 module(input) 的結果

返回型別

張量