如何告诉 Flux.jl 不使用 GPU?

How to tell Flux.jl not to use the GPU?

我正在尝试设置我的机器学习训练循环,但我不希望它使用我计算机上的 GPU。我已经设置并安装了 CUDA,但我想 运行 我的代码放在 CPU 上。我如何在 Flux.jl 中执行此操作?

默认情况下,除非您指定,否则 Flux 不会使用 GPU。首先,确保 CUDA 没有加载我的注释 using CUDA。接下来,删除对 gpu 的所有调用。一般来说,如果您希望在 CPU 和 GPU 上有选择地 运行 一个程序,那么使用以下设计模式编写您的代码:

using CUDA
CUDA.allowscalar(false)
using Flux

# configure device
device = gpu

# made up function that loads your dataset
data = loaddata() .|> device

# made up function that builds your model
m = mymodel() |> device

# rest of your flux code

请注意我是如何在各处使用通用 device 函数的,这样我只需要在程序启动时设置我选择的目标设备。如果您的程序在“main”函数中,您可以让该函数接受一个 device 参数。

如果您确实需要禁用 GPU,那么 CUDA.jl 会查找 CUDA_VISIBLE_DEVICES 环境变量,您可以将其设置为空。但是您只需要为调试 purposes/sanity 检查执行此操作。