如何在Julia Flux中为自定义Layer正确启用GPU支持？

我希望为自定义Layer添加GPU支持，以下是我的网络构建代码：

# MyDense
struct MyDense
    d1
    d2     
end
MyDense(p::Pair) = MyDense(Dense(p), Dense(p))
(m::MyDense)(x) = m.d1(x) + m.d2(x)
# multiple-dispatch for gpu function on the custom layer
Flux.gpu(layer::MyDense) = MyDense(Flux.gpu(layer.d1), Flux.gpu(layer.d2))

# Define our model, a multi-layer perceptron with one hidden layer of size 3:
model = Chain(
    Dense(2 => 3, tanh),   # activation function inside layer
    BatchNorm(3),
    MyDense(3 => 2),
    softmax) |> gpu        # move model to GPU, if available

我参考相关回答为自定义层编写了Flux.gpu的多分派方法，但运行时报错：

ERROR: ArgumentError: cannot take the CPU address of a CUDA.CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}
Stacktrace:
  [1] unsafe_convert(#unused#::Type{Ptr{Float32}}, x::CUDA.CuArray{Float32, 2, CUDA.Mem.DeviceBuffer})
    @ CUDA C:\Users\Herr LU\.julia\packages\CUDA\DfvRa\src\array.jl:319
  [2] gemm!(transA::Char, transB::Char, alpha::Float32, A::Matrix{Float32}, B::CUDA.CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, beta::Float32, C::CUDA.CuArray{Float32, 2, CUDA.Mem.DeviceBuffer})
    @ LinearAlgebra.BLAS C:\Users\Herr LU\AppData\Local\Programs\Julia-1.8.2\share\julia\stdlib\v1.8\LinearAlgebra\src\blas.jl:1514
  [3] gemm_wrapper!(C::CUDA.CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, tA::Char, tB::Char, A::Matrix{Float32}, B::CUDA.CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, _add::LinearAlgebra.MulAddMul{true, true, Bool, Bool})
    @ LinearAlgebra C:\Users\Herr LU\AppData\Local\Programs\Julia-1.8.2\share\julia\stdlib\v1.8\LinearAlgebra\src\matmul.jl:674
  [4] mul!
    @ C:\Users\Herr LU\AppData\Local\Programs\Julia-1.8.2\share\julia\stdlib\v1.8\LinearAlgebra\src\matmul.jl:161 [inlined]
  [5] mul!
    @ C:\Users\Herr LU\AppData\Local\Programs\Julia-1.8.2\share\julia\stdlib\v1.8\LinearAlgebra\src\matmul.jl:276 [inlined]
  [6] *
    @ C:\Users\Herr LU\AppData\Local\Programs\Julia-1.8.2\share\julia\stdlib\v1.8\LinearAlgebra\src\matmul.jl:148 [inlined]
  [7] (::Dense{typeof(identity), Matrix{Float32}, Vector{Float32}})(x::CUDA.CuArray{Float32, 2, CUDA.Mem.DeviceBuffer})
    @ Flux C:\Users\Herr LU\.julia\packages\Flux\ZdbJr\src\layers\basic.jl:172
  [8] (::MyDense)(x::CUDA.CuArray{Float32, 2, CUDA.Mem.DeviceBuffer})
    @ Main e:\Master Thesis\lu_jizhou\Learning\FluxTrial.jl:14
  [9] macro expansion
    @ C:\Users\Herr LU\.julia\packages\Flux\ZdbJr\src\layers\basic.jl:53 [inlined]
 [10] _applychain(layers::Tuple{Dense{typeof(tanh), CUDA.CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, CUDA.CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}, BatchNorm{typeof(identity), CUDA.CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}, Float32, CUDA.CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}, MyDense, typeof(softmax)}, x::CUDA.CuArray{Float32, 2, CUDA.Mem.DeviceBuffer})
    @ Flux C:\Users\Herr LU\.julia\packages\Flux\ZdbJr\src\layers\basic.jl:53
 [11] (::Chain{Tuple{Dense{typeof(tanh), CUDA.CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, CUDA.CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}, BatchNorm{typeof(identity), CUDA.CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}, Float32, CUDA.CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}, MyDense, typeof(softmax)}})(x::CUDA.CuArray{Float32, 2, CUDA.Mem.DeviceBuffer})
    @ Flux C:\Users\Herr LU\.julia\packages\Flux\ZdbJr\src\layers\basic.jl:51
 [12] top-level scope
    @ e:\Master Thesis\lu_jizhou\Learning\FluxTrial.jl:25

请问该如何解决这个问题？

问题根源 #

报错栈显示，MyDense内部的Dense层参数仍为CPU端的Matrix/Vector类型，但输入已经是GPU上的CuArray，导致CPU与GPU张量混合运算，触发了取CUDA数组CPU地址的错误。

问题出在MyDense的构造逻辑：调用MyDense(3 => 2)时生成的是CPU上的Dense层，后续的gpu方法虽尝试迁移，但需确保所有子层参数完全转移到GPU。

修复代码 #

修改自定义层的设备迁移逻辑，确保递归迁移所有子层，同时补充双向迁移的cpu方法：

# MyDense
struct MyDense
    d1
    d2     
end

MyDense(p::Pair) = MyDense(Dense(p), Dense(p))

# 递归迁移子层到GPU
Flux.gpu(layer::MyDense) = MyDense(Flux.gpu(layer.d1), Flux.gpu(layer.d2))
# 补充CPU迁移方法，保证双向兼容
Flux.cpu(layer::MyDense) = MyDense(Flux.cpu(layer.d1), Flux.cpu(layer.d2))

# 构建并迁移模型到GPU
model = Chain(
    Dense(2 => 3, tanh),
    BatchNorm(3),
    MyDense(3 => 2),
    softmax) |> gpu

# 验证子层参数是否已迁移到GPU
println(typeof(model[3].d1.weight))  # 应输出 CUDA.CuArray{Float32,2,...}

若仍有问题，可直接在构造MyDense时指定设备：

# 构造函数支持指定设备，默认使用GPU（若可用）
MyDense(p::Pair, device=gpu) = MyDense(device(Dense(p)), device(Dense(p)))

# 构建模型
model = Chain(
    Dense(2 => 3, tanh),
    BatchNorm(3),
    MyDense(3 => 2),
    softmax) |> gpu

核心注意事项 #

• 自定义层的gpu/cpu方法必须递归迁移所有子层与参数到对应设备
• 确保模型迁移后，所有参与运算的张量（层参数、输入）都在同一设备上
• 禁止CPU与GPU张量混合运算，这是触发此类错误的核心原因

内容的提问来源于stack exchange，提问作者PokeLu