site/ja/guide/advanced_autodiff.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Tce3stUlHN0L"
      },
      "source": [
        "##### Copyright 2020 The TensorFlow Authors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "tuOe1ymfHZPu"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qFdPvlXBOdUN"
      },
      "source": [
        "# 高度な自動微分"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MfBg1C5NB3X0"
      },
      "source": [
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td><a target=\"_blank\" href=\"https://www.tensorflow.org/guide/advanced_autodiff\"><img src=\"https://www.tensorflow.org/images/tf_logo_32px.png\">TensorFlow.org で表示</a></td>\n",
        "  <td>     <a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/docs-l10n/blob/master/site/ja/guide/advanced_autodiff.ipynb\">     <img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\">     Google Colab で実行</a>\n",
        "</td>\n",
        "  <td><a target=\"_blank\" href=\"https://github.com/tensorflow/docs-l10n/blob/master/site/ja/guide/advanced_autodiff.ipynb\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\">GitHub でソースを表示</a></td>\n",
        "  <td><a href=\"https://storage.googleapis.com/tensorflow_docs/docs-l10n/site/ja/guide/advanced_autodiff.ipynb\"><img src=\"https://www.tensorflow.org/images/download_logo_32px.png\">ノートブックをダウンロード</a></td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8a859404ce7e"
      },
      "source": [
        "[勾配と自動微分の基礎](autodiff.ipynb)には、TensorFlow での勾配の計算に必要なものがすべて含まれています。このガイドでは、`tf.GradientTape` API のさらに深い、あまり一般的ではない機能に焦点を当てていきます。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MUXex9ctTuDB"
      },
      "source": [
        "## セットアップ"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IqR2PQG4ZaZ0"
      },
      "outputs": [],
      "source": [
        "import tensorflow as tf\n",
        "\n",
        "import matplotlib as mpl\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "mpl.rcParams['figure.figsize'] = (8, 6)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uGRJJRi8TCkJ"
      },
      "source": [
        "## 勾配の記録を制御する\n",
        "\n",
        "[自動微分ガイド](autodiff.ipynb)では、勾配計算の作成中にテープが監視する対象の変数とテンソルを制御する方法について説明しました。\n",
        "\n",
        "テープには、記録を操作するメソッドもあります。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gB_i0VnhQKt2"
      },
      "source": [
        "### 記録を停止する\n",
        "\n",
        "勾配の記録を停止したい場合には、`tf.GradientTape.stop_recording`を使用して一時的に記録を停止することができます。\n",
        "\n",
        "これは、モデルの途中で複雑な演算を微分を行いたくない場合、オーバーヘッドの削減に有用です。これにはメトリックの計算や、中間結果の計算が含まれます。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mhFSYf7uQWxR"
      },
      "outputs": [],
      "source": [
        "x = tf.Variable(2.0)\n",
        "y = tf.Variable(3.0)\n",
        "\n",
        "with tf.GradientTape() as t:\n",
        "  x_sq = x * x\n",
        "  with t.stop_recording():\n",
        "    y_sq = y * y\n",
        "  z = x_sq + y_sq\n",
        "\n",
        "grad = t.gradient(z, {'x': x, 'y': y})\n",
        "\n",
        "print('dz/dx:', grad['x'])  # 2*x => 4\n",
        "print('dz/dy:', grad['y'])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DEHbEZ1h4p8A"
      },
      "source": [
        "### 記録をリセットまたは最初から録画し直す\n",
        "\n",
        "最初からやり直す場合は、`tf.GradientTape.reset` を使用してください。通常は勾配テープブロックを終了して再起動した方が読み取りやすくなりますが、テープブロックの終了が困難または不可能な場合は、`reset` メソッドを使用できます。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lsMHsmrh4pqM"
      },
      "outputs": [],
      "source": [
        "x = tf.Variable(2.0)\n",
        "y = tf.Variable(3.0)\n",
        "reset = True\n",
        "\n",
        "with tf.GradientTape() as t:\n",
        "  y_sq = y * y\n",
        "  if reset:\n",
        "    # Throw out all the tape recorded so far.\n",
        "    t.reset()\n",
        "  z = x * x + y_sq\n",
        "\n",
        "grad = t.gradient(z, {'x': x, 'y': y})\n",
        "\n",
        "print('dz/dx:', grad['x'])  # 2*x => 4\n",
        "print('dz/dy:', grad['y'])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6zS7cLmS6zMf"
      },
      "source": [
        "## 勾配フローを正確に停止する\n",
        "\n",
        "上記のグローバルなテープ制御とは対照的に、`tf.stop_gradient`関数ははるかに正確です。これはテープ自体にアクセスする必要がなく、勾配が特定のパスに沿って流れることを防ぐために使用できます。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "30qnZMe48BkB"
      },
      "outputs": [],
      "source": [
        "x = tf.Variable(2.0)\n",
        "y = tf.Variable(3.0)\n",
        "\n",
        "with tf.GradientTape() as t:\n",
        "  y_sq = y**2\n",
        "  z = x**2 + tf.stop_gradient(y_sq)\n",
        "\n",
        "grad = t.gradient(z, {'x': x, 'y': y})\n",
        "\n",
        "print('dz/dx:', grad['x'])  # 2*x => 4\n",
        "print('dz/dy:', grad['y'])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mbb-9lnGVngH"
      },
      "source": [
        "## カスタム勾配\n",
        "\n",
        "場合によっては、デフォルトを使用しないで勾配の計算方法を正確に制御する必要があります。次のような状況がそれに該当します。\n",
        "\n",
        "1. 記述中の新しい演算に勾配が定義されていない。\n",
        "2. デフォルトの計算が数値的に不安定。\n",
        "3. フォワードパスから高価な計算をキャッシュしたい。\n",
        "4. 勾配を変更せずに値を変更したい（`tf.clip_by_value` や `tf.math.round` などを使用）。\n",
        "\n",
        "最初のケースでは、新しい演算子を記述するために、`tf.RegisterGradient` を使用して独自の勾配を設定できます（詳細は API ドキュメントを参照）。（勾配レジストリはグローバルなので、注意して変更してください。）\n",
        "\n",
        "後者の 3 つのケースには、`tf.custom_gradient`を使用できます。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oHr31kc_irF_"
      },
      "source": [
        "以下は、中間勾配に `tf.clip_by_norm` を適用した例です。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Mjj01w4NYtwd"
      },
      "outputs": [],
      "source": [
        "# Establish an identity operation, but clip during the gradient pass.\n",
        "@tf.custom_gradient\n",
        "def clip_gradients(y):\n",
        "  def backward(dy):\n",
        "    return tf.clip_by_norm(dy, 0.5)\n",
        "  return y, backward\n",
        "\n",
        "v = tf.Variable(2.0)\n",
        "with tf.GradientTape() as t:\n",
        "  output = clip_gradients(v * v)\n",
        "print(t.gradient(output, v))  # calls \"backward\", which clips 4 to 2"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "n4t7S0scYrD3"
      },
      "source": [
        "詳細については、`tf.custom_gradient` デコレータの API ドキュメントをご覧ください。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "v0ODp4Oi--I0"
      },
      "source": [
        "### SavedModel のカスタム勾配\n",
        "\n",
        "注意: この機能は TensorFlow 2.6 以降で利用できます。\n",
        "\n",
        "カスタム勾配は、`tf.saved_model.SaveOptions(experimental_custom_gradients=True)` オプションを使用して、SavedModel に保存することができます。\n",
        "\n",
        "SavedModel に保存できるようにするには、勾配関数はトレース可能である必要があります（詳細は、「[tf.function によるパフォーマンスの改善](function.ipynb)」ガイドをご覧ください）。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Q5JBgIBYjN1I"
      },
      "outputs": [],
      "source": [
        "class MyModule(tf.Module):\n",
        "\n",
        "  @tf.function(input_signature=[tf.TensorSpec(None)])\n",
        "  def call_custom_grad(self, x):\n",
        "    return clip_gradients(x)\n",
        "\n",
        "model = MyModule()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xZTrgy2q-9pq"
      },
      "outputs": [],
      "source": [
        "tf.saved_model.save(\n",
        "    model,\n",
        "    'saved_model',\n",
        "    options=tf.saved_model.SaveOptions(experimental_custom_gradients=True))\n",
        "\n",
        "# The loaded gradients will be the same as the above example.\n",
        "v = tf.Variable(2.0)\n",
        "loaded = tf.saved_model.load('saved_model')\n",
        "with tf.GradientTape() as t:\n",
        "  output = loaded.call_custom_grad(v * v)\n",
        "print(t.gradient(output, v))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d-LfRs5FbJCk"
      },
      "source": [
        "上記の例に関する注意事項: 上記のコードを `tf.saved_model.SaveOptions(experimental_custom_gradients=False)` に置き換えても、勾配は読み込み時に同じ結果を生成します。これは、勾配のレジストリに `call_custom_op` 関数で使用されたカスタム勾配がまだ含まれているためです。ただし、カスタム勾配なしで保存した後にランタイムを再起動すると、読み込まれたモデルを `tf.GradientTape` で実行したときに `LookupError: No gradient defined for operation 'IdentityN' (op type: IdentityN)` というエラーがスローされます。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8aENEt6Veryb"
      },
      "source": [
        "## 複数のテープ\n",
        "\n",
        "複数のテープはシームレスに連携します。\n",
        "\n",
        "たとえば、ここでは各テープが異なるテンソルのセットをウォッチしています。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BJ0HdMvte0VZ"
      },
      "outputs": [],
      "source": [
        "x0 = tf.constant(0.0)\n",
        "x1 = tf.constant(0.0)\n",
        "\n",
        "with tf.GradientTape() as tape0, tf.GradientTape() as tape1:\n",
        "  tape0.watch(x0)\n",
        "  tape1.watch(x1)\n",
        "\n",
        "  y0 = tf.math.sin(x0)\n",
        "  y1 = tf.nn.sigmoid(x1)\n",
        "\n",
        "  y = y0 + y1\n",
        "\n",
        "  ys = tf.reduce_sum(y)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6ApAoMNFfNz6"
      },
      "outputs": [],
      "source": [
        "tape0.gradient(ys, x0).numpy()   # cos(x) => 1.0"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rF1jrAJsfYW_"
      },
      "outputs": [],
      "source": [
        "tape1.gradient(ys, x1).numpy()   # sigmoid(x1)*(1-sigmoid(x1)) => 0.25"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DK05KXrAAld3"
      },
      "source": [
        "### 高次勾配\n",
        "\n",
        "`tf.GradientTape` コンテキストマネージャー内の演算は、自動微分のために記録されます。そのコンテキストで勾配が計算される場合は、勾配の計算も記録されます。その結果、まったく同じ API が高次勾配にも機能します。\n",
        "\n",
        "以下に例を示します。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cPQgthZ7ugRJ"
      },
      "outputs": [],
      "source": [
        "x = tf.Variable(1.0)  # Create a Tensorflow variable initialized to 1.0\n",
        "\n",
        "with tf.GradientTape() as t2:\n",
        "  with tf.GradientTape() as t1:\n",
        "    y = x * x * x\n",
        "\n",
        "  # Compute the gradient inside the outer `t2` context manager\n",
        "  # which means the gradient computation is differentiable as well.\n",
        "  dy_dx = t1.gradient(y, x)\n",
        "d2y_dx2 = t2.gradient(dy_dx, x)\n",
        "\n",
        "print('dy_dx:', dy_dx.numpy())  # 3 * x**2 => 3.0\n",
        "print('d2y_dx2:', d2y_dx2.numpy())  # 6 * x => 6.0"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "k0HV-Ah4_76i"
      },
      "source": [
        "これにより、*スカラー*関数の 2 次導関数が得られますが、`tf.GradientTape.gradient` はスカラーの勾配のみを計算するため、このパターンはヘッセ行列を生成するための一般化を行いません。[ヘッセ行列](https://en.wikipedia.org/wiki/Hessian_matrix)を作成するには、[ヘッセ行列の例](#hessian)（[ヤコビアンのセクション](#jacobians)）をご覧ください。\n",
        "\n",
        "「`tf.GradientTape.gradient`のネストされた呼び出し」は、勾配からスカラーを計算する場合に適したパターンです。次の例のように、結果のスカラーは 2 番目の勾配計算のソースとして機能します。\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "t7LRlcpVKHv1"
      },
      "source": [
        "#### 例: 入力勾配の正則化\n",
        "\n",
        "多くのモデルは「敵対的な例」の影響を受けやすいものです。そのような手法の数々は、モデルの入力を変更してモデルの出力を混乱させます。[Fast Gradient Signed Method 攻撃を使った敵対的な例](https://www.tensorflow.org/tutorials/generative/adversarial_fgsm)といった最も単純な実装は、入力に対する出力の勾配に沿って 1 つのステップを踏む「入力勾配」です。\n",
        "\n",
        "敵対的な例に対する堅牢性を高める 1 つの手法は、入力勾配のマグニチュードの最小化を試みる、[入力勾配の正則化](https://arxiv.org/abs/1905.11468)（Finlay &amp; Oberman、2019 年）です。入力勾配が小さければ、出力の変化も小さくなるはずです。\n",
        "\n",
        "下記は、入力勾配の正則化の単純な実装です。実装は次のとおりです。\n",
        "\n",
        "1. 内部テープを使用して、入力に対する出力の勾配を計算する。\n",
        "2. その入力勾配の大きさを計算する。\n",
        "3. モデルに対する大きさの勾配を計算する。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tH3ZFuUfDLrR"
      },
      "outputs": [],
      "source": [
        "x = tf.random.normal([7, 5])\n",
        "\n",
        "layer = tf.keras.layers.Dense(10, activation=tf.nn.relu)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "E6yOFsjEDR9u"
      },
      "outputs": [],
      "source": [
        "with tf.GradientTape() as t2:\n",
        "  # The inner tape only takes the gradient with respect to the input,\n",
        "  # not the variables.\n",
        "  with tf.GradientTape(watch_accessed_variables=False) as t1:\n",
        "    t1.watch(x)\n",
        "    y = layer(x)\n",
        "    out = tf.reduce_sum(layer(x)**2)\n",
        "  # 1. Calculate the input gradient.\n",
        "  g1 = t1.gradient(out, x)\n",
        "  # 2. Calculate the magnitude of the input gradient.\n",
        "  g1_mag = tf.norm(g1)\n",
        "\n",
        "# 3. Calculate the gradient of the magnitude with respect to the model.\n",
        "dg1_mag = t2.gradient(g1_mag, layer.trainable_variables)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "123QMq6PqK_d"
      },
      "outputs": [],
      "source": [
        "[var.shape for var in dg1_mag]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "E4xiYigexMtQ"
      },
      "source": [
        "## ヤコビアン\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4-hVHVIeExkI"
      },
      "source": [
        "前述したすべての例では、複数または単一のソーステンソルに対し、単一のスカラーターゲットの勾配を取得しました。\n",
        "\n",
        "[ヤコビ行列](https://en.wikipedia.org/wiki/Jacobian_matrix_and_determinant)は、ベクトル値関数の勾配を表します。各行には、それぞれベクトルの一要素の勾配が示されます。\n",
        "\n",
        "`tf.GradientTape.jacobian`メソッドを使用すると、ヤコビ行列を効率的に計算することができます。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KzNyIM0QBYIH"
      },
      "source": [
        "注意点:\n",
        "\n",
        "- `gradient`と同様、`sources`引数は単一のテンソル、またはテンソルのコンテナである可能性があります。\n",
        "- `gradient`と異なり、`target`テンソルは単一のテンソルでなければなりません。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "O74K3hlxBC8a"
      },
      "source": [
        "### スカラーのソース"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "B08OKn1Orkuc"
      },
      "source": [
        "最初の例は、ここにスカラーソースに対するベクトルターゲットのヤコビアンです。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bAFeIE8EuVIq"
      },
      "outputs": [],
      "source": [
        "x = tf.linspace(-10.0, 10.0, 200+1)\n",
        "delta = tf.Variable(0.0)\n",
        "\n",
        "with tf.GradientTape() as tape:\n",
        "  y = tf.nn.sigmoid(x+delta)\n",
        "\n",
        "dy_dx = tape.jacobian(y, delta)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BgHbUk3zr-WU"
      },
      "source": [
        "スカラーに対してヤコビアンをとると、結果は**ターゲット**の形状を持ち、ソースに対する各要素の勾配を表します。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iZ6awnDzr_BA"
      },
      "outputs": [],
      "source": [
        "print(y.shape)\n",
        "print(dy_dx.shape)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "siNZaklc0_-e"
      },
      "outputs": [],
      "source": [
        "plt.plot(x.numpy(), y, label='y')\n",
        "plt.plot(x.numpy(), dy_dx, label='dy/dx')\n",
        "plt.legend()\n",
        "_ = plt.xlabel('x')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DsOMSD_1BGkD"
      },
      "source": [
        "### テンソルのソース"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "g3iXKN7KF-st"
      },
      "source": [
        "入力がスカラーとテンソルのどちらであっても、`tf.GradientTape.jacobian` は、単一または複数のターゲットの各要素に対するソースの各要素の勾配を効果的に計算します。\n",
        "\n",
        "たとえば、このレイヤーの出力の形状は `(10, 7)` です。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "39YXItgLxMBk"
      },
      "outputs": [],
      "source": [
        "x = tf.random.normal([7, 5])\n",
        "layer = tf.keras.layers.Dense(10, activation=tf.nn.relu)\n",
        "\n",
        "with tf.GradientTape(persistent=True) as tape:\n",
        "  y = layer(x)\n",
        "\n",
        "y.shape"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tshNRtfKuVP_"
      },
      "source": [
        "そしてレイヤーのカーネルの形状は`(5, 10)`です。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CigTWyfPvPuv"
      },
      "outputs": [],
      "source": [
        "layer.kernel.shape"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mN96JRpnAjpx"
      },
      "source": [
        "カーネルの出力のヤコビアンの形状は、次の 2 つの形状を連結したものです。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "pRLzTTbvEimH"
      },
      "outputs": [],
      "source": [
        "j = tape.jacobian(y, layer.kernel)\n",
        "j.shape"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2Lrv7miMvTll"
      },
      "source": [
        "ターゲットの次元を合計すると、`tf.GradientTape.gradient` によって計算された合計の勾配が残ります。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "FJjZpYRnDjVa"
      },
      "outputs": [],
      "source": [
        "g = tape.gradient(y, layer.kernel)\n",
        "print('g.shape:', g.shape)\n",
        "\n",
        "j_sum = tf.reduce_sum(j, axis=[0, 1])\n",
        "delta = tf.reduce_max(abs(g - j_sum)).numpy()\n",
        "assert delta < 1e-3\n",
        "print('delta:', delta)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZKajuGlk_krs"
      },
      "source": [
        "<a id=\"hessian\"> </a>\n",
        "\n",
        "#### 例: ヘッセ行列（Hessian）"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NYcsXeo8TDLi"
      },
      "source": [
        "`tf.GradientTape` には[ヘッセ行列](https://en.wikipedia.org/wiki/Hessian_matrix)を作成するための明示的なメソッドがありませんが、`tf.GradientTape.jacobian` を使用してこれを作成することは可能です。\n",
        "\n",
        "注意: ヘッセ行列には `N**2` 個のパラメータが含まれています。これやその他の理由により、ほとんどのモデルには非実用的です。あくまでも `GradientTape.jacobian` メソッドの使い方のデモンストレーションとしてこの例を含めましたが、直接のヘッセ行列ベースの最適化を推奨するものではありません。ヘッセ行列のベクトル積は[ネストされたテープで効率的に計算できる](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/eager/benchmarks/resnet50/hvp_test.py)ため、2 次最適化の方法としては、はるかに効率的です。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ELGTaell_j81"
      },
      "outputs": [],
      "source": [
        "x = tf.random.normal([7, 5])\n",
        "layer1 = tf.keras.layers.Dense(8, activation=tf.nn.relu)\n",
        "layer2 = tf.keras.layers.Dense(6, activation=tf.nn.relu)\n",
        "\n",
        "with tf.GradientTape() as t2:\n",
        "  with tf.GradientTape() as t1:\n",
        "    x = layer1(x)\n",
        "    x = layer2(x)\n",
        "    loss = tf.reduce_mean(x**2)\n",
        "\n",
        "  g = t1.gradient(loss, layer1.kernel)\n",
        "\n",
        "h = t2.jacobian(g, layer1.kernel)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "FVqQuZj4XGjm"
      },
      "outputs": [],
      "source": [
        "print(f'layer.kernel.shape: {layer1.kernel.shape}')\n",
        "print(f'h.shape: {h.shape}')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_M7XElgaiMeP"
      },
      "source": [
        "このヘッセ行列をニュートン法のステップに使用するには、まず軸を平坦化して行列にし、勾配を平坦化してベクトルにします。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6te7N6wVXwXX"
      },
      "outputs": [],
      "source": [
        "n_params = tf.reduce_prod(layer1.kernel.shape)\n",
        "\n",
        "g_vec = tf.reshape(g, [n_params, 1])\n",
        "h_mat = tf.reshape(h, [n_params, n_params])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "L9rO8b-0mgOH"
      },
      "source": [
        "ヘッセ行列は必ず対称行列になります。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8TCHc7Vrf52S"
      },
      "outputs": [],
      "source": [
        "def imshow_zero_center(image, **kwargs):\n",
        "  lim = tf.reduce_max(abs(image))\n",
        "  plt.imshow(image, vmin=-lim, vmax=lim, cmap='seismic', **kwargs)\n",
        "  plt.colorbar()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DExOxd7Ok2H0"
      },
      "outputs": [],
      "source": [
        "imshow_zero_center(h_mat)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "13fBswmtQes4"
      },
      "source": [
        "ニュートン法の更新ステップを以下に示します。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3DdnbynBdSor"
      },
      "outputs": [],
      "source": [
        "eps = 1e-3\n",
        "eye_eps = tf.eye(h_mat.shape[0])*eps"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-zPdtyoWeUeV"
      },
      "source": [
        "注意: [実際に行列を反転させないでください](https://www.johndcook.com/blog/2010/01/19/dont-invert-that-matrix/)。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "k1LYftgmswOO"
      },
      "outputs": [],
      "source": [
        "# X(k+1) = X(k) - (∇²f(X(k)))^-1 @ ∇f(X(k))\n",
        "# h_mat = ∇²f(X(k))\n",
        "# g_vec = ∇f(X(k))\n",
        "update = tf.linalg.solve(h_mat + eye_eps, g_vec)\n",
        "\n",
        "# Reshape the update and apply it to the variable.\n",
        "_ = layer1.kernel.assign_sub(tf.reshape(update, layer1.kernel.shape))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pF6qjlHKWxF4"
      },
      "source": [
        "これは、単一の`tf.Variable`の場合には比較的簡単ですが、これを複雑なモデルに適用するには、完全なヘッセ行列を複数の変数にわたって生成するために、注意深い連結とスライスが必要です。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PQWM0uN-GO5t"
      },
      "source": [
        "### ヤコビアンをバッチ処理する"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hKtB3rY6EySJ"
      },
      "source": [
        "場合により、複数のソースの単一スタックに対して、複数のターゲットの各スタックのヤコビアンを取得して、ターゲットとソースの各ペアのヤコビアンを独立させる必要があることがあります。\n",
        "\n",
        "たとえば、ここでは入力 `x` は `(batch, ins)` の形状をしており、出力 `y` は `(batch, outs)` の形状をしています。\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tQMndhIUHMes"
      },
      "outputs": [],
      "source": [
        "x = tf.random.normal([7, 5])\n",
        "\n",
        "layer1 = tf.keras.layers.Dense(8, activation=tf.nn.elu)\n",
        "layer2 = tf.keras.layers.Dense(6, activation=tf.nn.elu)\n",
        "\n",
        "with tf.GradientTape(persistent=True, watch_accessed_variables=False) as tape:\n",
        "  tape.watch(x)\n",
        "  y = layer1(x)\n",
        "  y = layer2(y)\n",
        "\n",
        "y.shape"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ff2spRHEJXBU"
      },
      "source": [
        "`x` に対する `y` の完全なヤコビアンは、たとえ `(batch, ins, outs)` だけが必要な場合でも、`(batch, ins, batch, outs)` の形状をしています。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1zSl2A5-HhMH"
      },
      "outputs": [],
      "source": [
        "j = tape.jacobian(y, x)\n",
        "j.shape"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UibJijPLJrpQ"
      },
      "source": [
        "スタック内の各アイテムの勾配が独立している場合は、このテンソルの`(batch, batch)`スライスはすべて対角行列になります。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZFl9uj3ueVSH"
      },
      "outputs": [],
      "source": [
        "imshow_zero_center(j[:, 0, :, 0])\n",
        "_ = plt.title('A (batch, batch) slice')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "g4ZoRJcJNmy5"
      },
      "outputs": [],
      "source": [
        "def plot_as_patches(j):\n",
        "  # Reorder axes so the diagonals will each form a contiguous patch.\n",
        "  j = tf.transpose(j, [1, 0, 3, 2])\n",
        "  # Pad in between each patch.\n",
        "  lim = tf.reduce_max(abs(j))\n",
        "  j = tf.pad(j, [[0, 0], [1, 1], [0, 0], [1, 1]],\n",
        "             constant_values=-lim)\n",
        "  # Reshape to form a single image.\n",
        "  s = j.shape\n",
        "  j = tf.reshape(j, [s[0]*s[1], s[2]*s[3]])\n",
        "  imshow_zero_center(j, extent=[-0.5, s[2]-0.5, s[0]-0.5, -0.5])\n",
        "\n",
        "plot_as_patches(j)\n",
        "_ = plt.title('All (batch, batch) slices are diagonal')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OXpTBKyeK84z"
      },
      "source": [
        "期待どおりの結果を得るためには、重複する `batch` 次元を合計するか、`tf.einsum` を使用して対角を選択します。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "v65OAjEgLQwl"
      },
      "outputs": [],
      "source": [
        "j_sum = tf.reduce_sum(j, axis=2)\n",
        "print(j_sum.shape)\n",
        "j_select = tf.einsum('bxby->bxy', j)\n",
        "print(j_select.shape)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zT_VfR6lcwxD"
      },
      "source": [
        "そもそも、余分な次元なしで計算を行う方がはるかに効率的です。`tf.GradientTape.batch_jacobian` メソッドは、まさにそれを行います。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YJLIl9WpHqYq"
      },
      "outputs": [],
      "source": [
        "jb = tape.batch_jacobian(y, x)\n",
        "jb.shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-5t_q5SfHw7T"
      },
      "outputs": [],
      "source": [
        "error = tf.reduce_max(abs(jb - j_sum))\n",
        "assert error < 1e-3\n",
        "print(error.numpy())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IUeY2ZCiL31I"
      },
      "source": [
        "注意: `tf.GradientTape.batch_jacobian` は、ソースとターゲットの最初の次元が一致することのみを検証します。勾配が実際に独立していることは確認しません。意味のある場所でのみ `batch_jacobian` を使用できるかは、ユーザー次第です。たとえば、`tf.keras.layers.BatchNormalization` を追加すると、`batch` 次元の全体にわたって正規化するため、独立性は破壊されます。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tnDugVc-L4fj"
      },
      "outputs": [],
      "source": [
        "x = tf.random.normal([7, 5])\n",
        "\n",
        "layer1 = tf.keras.layers.Dense(8, activation=tf.nn.elu)\n",
        "bn = tf.keras.layers.BatchNormalization()\n",
        "layer2 = tf.keras.layers.Dense(6, activation=tf.nn.elu)\n",
        "\n",
        "with tf.GradientTape(persistent=True, watch_accessed_variables=False) as tape:\n",
        "  tape.watch(x)\n",
        "  y = layer1(x)\n",
        "  y = bn(y, training=True)\n",
        "  y = layer2(y)\n",
        "\n",
        "j = tape.jacobian(y, x)\n",
        "print(f'j.shape: {j.shape}')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SNyZ1WhJMVLm"
      },
      "outputs": [],
      "source": [
        "plot_as_patches(j)\n",
        "\n",
        "_ = plt.title('These slices are not diagonal')\n",
        "_ = plt.xlabel(\"Don't use `batch_jacobian`\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "M_x7ih5sarvG"
      },
      "source": [
        "この場合、`batch_jacobian` は引き続き実行され、期待される形状の*何か*を返しますが、その内容の意味は不明確です。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "k8_mICHoasCi"
      },
      "outputs": [],
      "source": [
        "jb = tape.batch_jacobian(y, x)\n",
        "print(f'jb.shape: {jb.shape}')"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "advanced_autodiff.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}