TensorFlowのOptimizerを自作する

機械学習

TensorFlowのOptimizerを自分で独自に実装することができるようなので、実装してみたソースコードをメモとして残しておきます。以下の私の記事で使ったアルゴリズムを実装してみました。これで自分のアルゴリズムをTensorFlowで動かせます。

以下の記事ではこのOptimizerを実際に使っています。

これらの記事のPythonのソースコード中に CustomOptimizer というのがありますが、全部以下のソースコードで定義しています。

import numpy as np
import math
import tensorflow as tf

# オリジナルのOptimizer
class CustomOptimizer(tf.optimizers.Optimizer):
  def __init__(self, learning_rate=0.1, name='CustomOptimizer', **kwargs):
    super(CustomOptimizer, self).__init__(name, **kwargs)
    self.learning_rate = kwargs.get('learning_rate', learning_rate)

  def get_config(self):
    config = super(CustomOptimizer, self).get_config()
    config.update({
        'learning_rate': self.learning_rate
    })
    return config
  
  def _create_slots(self, var_list):
    for v in var_list:
      self.add_slot(v, 'grad', tf.zeros_like(v))
      self.add_slot(v, 'diff', tf.zeros_like(v))

  def _resource_apply_dense(self, grad, var):
    k = math.log(3.0)
    prev_grad = self.get_slot(var, 'grad').numpy()
    prev_diff = self.get_slot(var, 'diff').numpy()
    grad = grad.numpy()
    diff = np.copy(grad * self.learning_rate)
    replace_mask = (prev_grad != 0.0)
    diff[replace_mask] = prev_diff[replace_mask] * (0.5 + np.tanh(k * (1.5 * grad[replace_mask] / prev_grad[replace_mask] - 0.5)))
    self.get_slot(var, 'grad').assign(grad)
    self.get_slot(var, 'diff').assign(diff)
    return var.assign_sub(diff)

Google Colaboratoryで動かしていますが、TensorFlowのバージョンは2.3.0です。