Juliaの自動微分パッケージ Zygote の紹介

Zygoteパッケージ

Juliaの数値微分パッケージであるZygoteパッケージを簡単に紹介します。
Julia/Zygoteでは以前取り扱ったPyTorchの自動微分に比べ非常に簡単に少ないコード量で同等のことが実行できます。

Table of Contents

1. Zygoteパッケージ
2. 公式ページ
3. Setup
   1. 環境
   2. インストール
4. 簡単な自動微分の実行
   1. Julia
   2. PyTorch
5. その他 Zygote の利用方法

公式ページ

github : https://github.com/FluxML/Zygote.jl
docs : https://fluxml.ai/Zygote.jl/dev/

Setup

環境

• julia version 1.6.0
• Zygote v0.6.8

インストール

Zygoteは、JuliaのREPLのパッケージマネージャーからインストールできます。

pkg> Zygote

簡単な自動微分の実行

今回紹介するパッケージを用いてJuliaとPyTorchで比較します。

Julia

まずはパッケージをインポートします

using Zygote
using Plots

微分したい関数と参考点を用意します。

f(x) = sin(x);
inputs = range(0, 4pi, length=50)

元の関数を可視化します。

plot(inputs, f.(inputs), label="f(x)=sin(x)")

微分の実行と、微分した関数の可視化を同時に行います。

plot!(inputs, f'.(inputs), label="f'(x)=cos(x)")
plot!(inputs, f''.(inputs), label="f''(x)=-sin(x)")

微分の実行は非常に簡単です。 ' だけで微分が実行されます。また、 ' を重ねればn階微分も可能です。

PyTorch

まずはパッケージをインポートします

import torch
import numpy as np
import matplotlib.pyplot as plt

微分したい関数と参考点を用意します。

f = lambda x: torch.sin(x)
inputs = torch.linspace(0, 4*np.pi, 50, requires_grad=True)

元の関数を可視化します。

plt.plot(inputs.detach().numpy(), f(inputs).detach().numpy(), label="f(x)=sin(x)")

微分の実行と可視化を同時に行います。

y = f(inputs).sum()
grad = torch.autograd.grad(outputs=y, inputs=inputs, create_graph=True)
grad[0].sum().backward()
plt.plot(inputs.detach().numpy(), grad[0].detach().numpy(), label="f'(x)=cos(x)")
plt.plot(inputs.detach().numpy(), inputs.grad.detach().numpy(), label="f''(x)=-sin(x)")
plt.grid()
plt.legend()
plt.show()

PyTorchの微分メソッドと比べると、プライムで微分ができるのは慣れてない人にも扱いやすいですね。

その他 Zygote の利用方法

上記のように単純な微分だけでも遊びの幅は広がりますが、より機械学習としての実用という観点で様々な実装がされているようです。
特に行列の微分( gradient )は利用機会は多いと思いますが下記の公式ドキュメントから参照してください。
docs : https://fluxml.ai/Zygote.jl/dev/