Question

我在.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))

Answer 1

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
}

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

1 个答案: