Tensorflow 2.Xでの線形回帰

今回はTensorflow2で線形回帰を行ってみようと思います．

Linear Regression using TensorFlow 2.0 | by Dhiraj K | Heartbeat
のページを参考にTensorflow 2を用いた線形回帰のメモを残そうと思います．
※一部改変しています

インポート

import tensorflow as tf
import numpy as np

回帰用データ作成

x_train = np.linspace(-2, 2, 400)
y_train = 3*x_train**2 + 5*x_train + 3 + np.random.randn(400)

xを400点，yは適当な二次関数にノイズを加えたものを採用しました．
この二次の係数3，一次の係数5，定数項3をモデルが学習し，"それっぽい"線を引くことが目的です．
今回は面倒なのでテストデータは用意していません．

モデル定義とインスタンス化

class MyModel(tf.keras.Model):
  def __init__(self):
    super(MyModel, self).__init__()

    self.a1 = tf.Variable(1.0)
    self.a2 = tf.Variable(1.0)
    self.b = tf.Variable(1.0)
    self.params = [self.a1, self.a2, self.b]

  def call(self, x):
    return self.a1*x**2 + self.a2*x + self.b

model = MyModel()

2系から学び始めたため，VariableはPlaceholder等と共に既に廃止されたものと勝手に思っていましたが，普通に使えるのですね．
複数の層を用いたMNISTの分類では，call関数でフィードフォワードの繋がりを決定し出力である確率を返しましたが，今回の出力は単純に係数を掛けた予測値です．
参考記事のおかげで，call関数は入力から出力のプロセスを記述し出力を返す関数であると再認識できました．

損失関数と最適化手法の定義

loss_object = tf.keras.losses.MeanSquaredError()
optimizer = tf.keras.optimizers.SGD(learning_rate=0.01)

参考記事では，以下のようなloss関数を自ら定義しています．

def loss(y, pred):
    return tf.reduce_mean(tf.square(y - pred))

ただ，これってMeanSquaredErrorと同義じゃね？と思ったのでtf.keras.losses.MeanSquaredError()を用いました．

オプティマイザに関しても，参考記事では，

lr_weight, lr_bias = t.gradient(current_loss, [linear_model.Weight, linear_model.Bias])
linear_model.Weight.assign_sub(lr * lr_weight)
linear_model.Bias.assign_sub(lr * lr_bias)

のように，lossとモデルの各変数をgradientに与えて勾配を求め，assign_subというメソッドを使って，勾配と学習係数に掛け合わせた値で各変数を更新しているみたいですが，おそらくSGDと同じ原理ですので，今回はtf.keras.optimizers.SGD(learning_rate=0.01)としました．

このassign_subについて調べて見ると1系の関数で2系には非対応っぽいんですがどうなんでしょう…

学習用関数定義

@tf.function
def train_step(x, y):
  with tf.GradientTape() as tape:
    predictions = model(x)
    loss = loss_object(y, predictions)
  
  gradients = tape.gradient(loss, model.params)
  optimizer.apply_gradients(zip(gradients, model.trainable_variables))

1. 予測値(モデルの出力)を求める
2. 正解と予測値からlossを求める
3. lossとモデルの各変数から勾配を求める
4. 勾配を用いてSGDで重みを更新

model.paramsはモデルの各変数をリスト形式にまとめたものです．

学習

EPOCH = 100

for epoch in range(EPOCH):
  # 学習
  train_step(x_train, y_train)

  # 正解と予測値からlossを求める（出力用）
  real_loss = loss_object(y_train, model(x_train))

  # 出力
  template = "Epoch {}, Loss {}"
  print(template.format(epoch+1, real_loss.numpy()))

特筆すべきことはありません．

結果

Epoch 1, Loss 45.20094680786133
Epoch 2, Loss 40.685359954833984
Epoch 3, Loss 36.73365020751953
Epoch 4, Loss 33.267723083496094
Epoch 5, Loss 30.220775604248047
Epoch 6, Loss 27.53564453125
Epoch 7, Loss 25.163389205932617
Epoch 8, Loss 23.062070846557617
Epoch 9, Loss 21.195751190185547
Epoch 10, Loss 19.53359603881836
Epoch 11, Loss 18.04913330078125
Epoch 12, Loss 16.71961212158203
Epoch 13, Loss 15.525479316711426
Epoch 14, Loss 14.4498929977417
Epoch 15, Loss 13.478344917297363
Epoch 16, Loss 12.598313331604004
Epoch 17, Loss 11.798981666564941
Epoch 18, Loss 11.070996284484863
Epoch 19, Loss 10.40625286102295
Epoch 20, Loss 9.797710418701172
Epoch 21, Loss 9.239261627197266
Epoch 22, Loss 8.725581169128418
Epoch 23, Loss 8.252019882202148
Epoch 24, Loss 7.814520835876465
Epoch 25, Loss 7.409524917602539
Epoch 26, Loss 7.033907413482666
Epoch 27, Loss 6.684916973114014
Epoch 28, Loss 6.3601250648498535
Epoch 29, Loss 6.0573859214782715
Epoch 30, Loss 5.774795055389404
Epoch 31, Loss 5.510659217834473
Epoch 32, Loss 5.263465404510498
Epoch 33, Loss 5.031863212585449
Epoch 34, Loss 4.814640045166016
Epoch 35, Loss 4.61070442199707
Epoch 36, Loss 4.419074535369873
Epoch 37, Loss 4.238860130310059
Epoch 38, Loss 4.0692548751831055
Epoch 39, Loss 3.9095232486724854
Epoch 40, Loss 3.758997917175293
…（中略）…
Epoch 90, Loss 1.2808070182800293
Epoch 91, Loss 1.2708830833435059
Epoch 92, Loss 1.2614418268203735
Epoch 93, Loss 1.252458095550537
Epoch 94, Loss 1.2439079284667969
Epoch 95, Loss 1.235769271850586
Epoch 96, Loss 1.2280209064483643
Epoch 97, Loss 1.2206425666809082
Epoch 98, Loss 1.2136154174804688
Epoch 99, Loss 1.2069215774536133
Epoch 100, Loss 1.2005436420440674

参考記事では一次関数の近似でlossが20エポックほどで1を下回っているんですが，今回は二次関数を利用したからなのか，はたまた改変した部分が影響を及ぼしたのかわかりませんが，100エポックでlossは1程度という結果になりました．

各係数を表示

# モデルの重みをリスト形式で取得
model.get_weights() # [3.1935437, 4.79401, 2.541419]

# また，以下のように直接モデルの変数を参照してもよいです
model.a1.numpy() # 3.1935437
model.a2.numpy() # 4.79401
model.b.numpy() # 2.541419

理想はそれぞれa1=3,a2=5,b=3ですので，バイアス項のノイズを考慮すればこんなものなのかもしれません（適当）．

最後に，得られた係数で描画した曲線を，元のグラフと重ねて貼っておきます．

個人的に，各係数の値だけを理想値と比較するとズレが大きいように感じていましたが，グラフにプロットしてみるとかなり忠実に元グラフをなぞれていますね！
線形回帰の場合，必ずしも各係数を完璧なレベルで近似する必要はないのかもしれません．