julia - 如何告诉 Flux.jl 不要使用 GPU？

Question

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

Answer 1

默认情况下，除非您告诉它，否则 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您可以设置为空的环境变量。但是您只需要为调试目的/完整性检查执行此操作。