时间戳

首页 > 解决方案 > 如何看Zygote差异化功能实现？

julia - 如何看Zygote差异化功能实现？

问题描述

我在一个.jl文件中编写了一个简单的函数，我可以使用forward. 但是我是 Julia 的新手，我不明白如何查看生成的微分函数的源代码。我已经尝试了各种各样的事情@code_lowered Zygote.forward(maxPool, [1.0, 2.0])，@code_lowered Zygote.forward(maxPool)但他们只是向我展示了转发本身的呼叫。

如何查看 Zygote 为正向和反向传递生成的代码？

using Pkg
using Zygote, ForwardDiff

function size1d(v)
    return size(v)[1]
end

function max(a, b)
    if a > b
        a
    else
        b
    end
end

function maxPool(v)
    return [max(v[2 * i - 1], v[2 * i])
            for i in 1:div(size1d(v), 2)]
end

v = [1.0, 2.0, 3.0, 4.0]
df = [20.0, 30.0]

println("maxPool(v):")
println(maxPool(v))
println()

println("maxAdjoint:")
maxAdjoint = Zygote.forward(max, 3.0, 4.0)[2]
println(maxAdjoint(1.0))
println()

println("maxPoolAdjoint:")
maxPoolAdjoint = Zygote.forward(maxPool, v)[2]
println(maxPoolAdjoint(df))

标签： julia

解决方案

Zygote 有自己的宏Zygote.@show_adjoint用于显示降低的伴随代码，即在反向模式下生成函数梯度的代码。不过，我不确定前进模式。

下面是一个反向模式的简单示例：

julia> using Zygote

julia> f(x) = 2x + 1
f (generic function with 1 method)

julia> @code_lowered f(1)
CodeInfo(
1 ─ %1 = 2 * x
│   %2 = %1 + 1
└──      return %2
)

julia> Zygote.@code_adjoint f(1)
Zygote.Adjoint(1: (%3, %4 :: Zygote.Context, %1, %2)
  %5 = Zygote._forward(%4, Main.:*, 2, %2)
  %6 = Base.getindex(%5, 1)
  %7 = Base.getindex(%5, 2)
  %8 = Zygote._forward(%4, Main.:+, %6, 1)
  %9 = Base.getindex(%8, 1)
  %10 = Base.getindex(%8, 2)
  return %9
, 1: (%1)
  %2 = (@10)(%1)
  %3 = Zygote.gradindex(%2, 2)
  %4 = (@7)(%3)
  %5 = Zygote.gradindex(%4, 3)
  %6 = Zygote.tuple(nothing, %5)
  return %6
)

我们可能会担心这种降低的伴随代码的长度和明显的复杂性会导致梯度变慢，但我们可以检查 LLVM 代码以确保所有内容最终都被忽略：

julia> @code_llvm f'(1)

;  @ /Users/mason/.julia/packages/Zygote/SAZMM/src/compiler/interface.jl:50 within `#34'
define i64 @"julia_#34_18250"(i64) {
top:
  ret i64 2
}

推荐阅读