Colab TPUでもtensorflow.kerasでBilinear法のアップサンプリングを行う方法

前置き

Google Colaboratoryでは、無料でクッソ早いプロセッサであるTPUを使って機械学習ができます。（Google ColaboratoryのTPUについて、詳しくは「Google ColaboratoryのTPUを試してみる」という記事を見ればわかりやすいです。）
しかしTPUには制限があり、使える関数が限られています。例えばkeras.layersのUpsampling2DはTPUでは使用ません。（対処法に関しては「TPUでアップサンプリングする際にエラーを出さない方法」が詳しいです。）
TPUで使える関数についての詳細は「利用可能な TensorFlow オペレーション」で分かるのですが内容がtensorflow 2.xに対応してないので、tensorflow 2.xのtf.image.resizeでBilinear法のアップサンプリングをしたら予期せぬエラーや動作にハマって地獄を見ることになるかもしれないのです。
そこで今回はtensorflow 2.xでもBilinear法のアップサンプリングができる方法を記したいと思います。

実装

tensorflow 2.xでもtensorflow 1.xの関数を使うにはtf.compat.v1から目的の関数を呼び出せばいいので自ら実装はしません。（下手に実装したら遅くなるだろうし...）
TPUに対応していることが判明しているtensorflow 1.xのBilinear法のアップサンプリングの関数は、tf.image.resize_bilinearで、変数をalign_corners=Trueにしたときのみ対応しています。
よって、tf.compat.v1.image.resize_bilinearを使います。
この関数をkerasのレイヤーにすると以下のようになります。

from tensorflow.keras.layers import *

# Lambdaレイヤーを使う
def upsampling_tpu_bilinear(x, scale=2):
    w = x.shape[1]*scale; h = x.shape[2]*scale
    tf.compat.v1.image.resize_bilinear(x, (w, h), align_corners=True)

Lambda(upsampling_tpu_bilinear, arguments={'scale': 2})(inputs)

最後に

これでTPUでもBilinear法のアップサンプリングができます。
「モザイク除去から学ぶ　最先端のディープラーニング」という本（めちゃくちゃオススメです！）でも解説されていますが、Nearest Neighbor法でアップサンプリングされた画像は「おかしな」画像なので、GANのような繊細なネットワークではBilinear法のアップサンプリングも効果的かもしれませんね。