JuliaのITensorsでGPUとCPUを行き来する

ITensorsとGPU

GPUを用いると、DMRGなどの計算が早くなるため非常に快適です。

GPUとCPUを変換

実際に変換してみましょう。次のように準備しておきます。

using CUDA
using ITensors, ITensorMPS

const L = 10
const sites = siteinds("S=1/2",L)

ここで、MPSを用意しておきましょう。（これは、CPU側にあります）

ψ = random_mps(sites,linkdims=2)

CPU -> GPU

こちらのサイトに従って、CPUからGPUに転送してみます。

ψ_gpu = NDTensors.CUDAExtensions.cu(ψ)

あるいは、ただ単に

ψ_gpu = cu(ψ)

とします。
これだけでCPUからGPUへと転送することができます。

GPU -> CPU

こちらのサイトに従って、GPUからCPUに転送してみます。

ψ_cpu = NDTensors.cpu(ψ_gpu)

これだけでGPUからCPUへと転送することができます。

チェック

実際に、ψとψ_cpuが同一であるか確認してみます。

julia> ψ ≈ ψ_cpu
true

但し、CPUにあるψでは倍精度で定義されますが、cu()で転送したψ_gpuでは単精度になってしまう影響によるものか、厳密には一致しないので注意です。

julia> ψ == ψ_cpu
false

最後に

CPU -> GPUについての転送は記事が多くあったのですが、
GPU -> CPUについて、日本語の記事が見つからなかったため、作成しました。
また、CPU -> GPUへの転送の際に、Float64のGPUへと転送できる方法を探していますので、ご存じの方はコメントいただけると幸いです。

追記

Float64で変換するために、一旦Arrayに焼き直したのちに、CuArrayでGPUへと倍精度で送り、それをITensorで焼き直す、といった方法はとれそうです。

function mps_to_gpu(ψ::MPS)
    s = siteinds(ψ)
    ψ = orthogonalize!(ψ, 1)
    A = ψ[1]
    for j = 2:length(ψ)
        A = A * ψ[j]
    end
    A_array = Array(A, s...)
    return MPS(CuArray(A_array),s)
end

例えば、次のように用います。

ψ_gpu_ = mps_to_gpu(ψ)

これで、倍精度にすることができました。

但し、こちらの場合でも、もともとのとは一致しませんでした。

julia> ψ == NDTensors.cpu(ψ_gpu_)
false