TPUを仕組み理解からTensorFlowで実装まで

2021年4月3日

TPUって何？

見ていると気になることはいっぱいある
上げればキリがないものだが、昔のメモをさグッたらTPUという文字が出てきたのでこの際に調べておこう。

Tensor Processing Unit

最も事細かにかかなので、あしからず。
仕組みの説明はGoogleCloudの説明がしっくりきた。
まず、CPU・GPUの復習

CPU
- ALUに何でもやらせるから柔軟性は高い。
- でもキャッシュ読み込み→命令処理→書き込みまでが時間とエネルギーの無駄
- パイプライン処理で頑張っても、結局たかが知れてる
GPU
- 1個のプロセッサに大量のALUを入れたもの。単純な命令を「数」で倒せる。
- 数の分だけ高速化するが、電力を食う。
- CPUと仕組みは同じなので、中間でキャッシュを経由するため時間とエネルギーの無駄
  - この呪縛をノイマンボトルネックというらしい。この名は初めて知った

Googleはうまいことを考えた。GPU・CPUの欠点を克服するには、構造をかえようと
キャッシュに行く前にどんどん足してしまえばよい！
Google公式のGIFがよいね

データを左から列方向に展開して→に進めていく

パラメータを下から行方向に展開して↑に進めていく

NNっていうのは入力データに対し、パラメータを1対1(y=ax+b)でつなげるものなので、
上２図のように流していけば、否が応でも演算器にぶち当たり計算せざるを得ない。
出口である右端では最終的に下みたいな行列演算(matmul)ができていることになる

感動した。頭いい。

このキャッシュ通さない作戦をシストリックアレイって言うらしい。
また、NNの計算過程では精度はそこまで重要ではないらしく、
量子化を行って一般的な32/64bit演算から8bit演算の仕組みを採用したため、実質４倍高速化(=省エネ)
この精度周りの下りは結構アツいらしい
https://ascii.jp/elem/000/004/014/4014066/

ムーアの法則って呪縛に聞こえてきた…

GPUとCPUのお互いいいところがあるように、TPUにも向き不向きがある。
Googleから直接拝借する。

・CPU
　最大限の柔軟性を必要とする迅速なプロトタイピング
　トレーニングに時間がかからない単純なモデル
　実際のバッチサイズが小さい小規模なモデル
　C++ で記述されたカスタム TensorFlow 演算が多くを占めるモデル
　ホストシステムの使用可能な I/O またはネットワーク帯域幅によって制限が課せられるモデル
・GPU
　ソースが存在しないモデルまたはソースを変更するのが煩雑すぎるモデル
　CPU 上で少なくとも部分的に実行しなければならない多数のカスタム TensorFlow 演算を使用するモデル
　Cloud TPU で利用できない TensorFlow 演算を使用するモデル（利用可能な TensorFlow 演算のリストをご覧ください）
　実際のバッチサイズが大きい中〜大規模なモデル
・TPU
　行列計算が多くを占めるモデル
　メインのトレーニングループ内にカスタム TensorFlow 演算がないモデル
　トレーニングに数週間または数か月かかるモデル
　実際のバッチサイズが非常に大きい非常に大規模なモデル
https://cloud.google.com/tpu/docs/tpus?hl=ja

GPUはある程度複雑な計算でもある程度応用(適応？)が効いたが、TPUは一切それが効かない。
その代わり、単純な行列演算の塊なら必殺！といった次第だろう

ColabとかKaggleとかでは使えるのでつかってみた

実装してみます。とは言っても他人コードを拝借する。一部加えつつ、、
１年前にkaggaleでTPUつかおうぜ！みたいなのを思い出したので、それを題材にします。
お花分類ですね

こちらに書いてあるTPUの使い方はとてもわかりやすいと思います（あくまで理解のうえでは）
SAMPLEコードにめもを追加していきます。

import re
import tensorflow as tf
import tensorflow_hub as hub
import numpy as np
from matplotlib import pyplot as plt
print("Tensorflow version " + tf.__version__)

# Prefetch用にどの程度先取りしておくか自動調整
#　これもいつもどおりでOK
# https://tensorflow.classcat.com/2019/03/23/tf20-alpha-guide-data-performance/
AUTO = tf.data.experimental.AUTOTUNE
# Kaggleのデータ拝借
from kaggle_datasets import KaggleDatasets

import re

import tensorflow as tf

import tensorflow_hub as hub

import numpy as np

from matplotlib import pyplot as plt

print("Tensorflow version " + tf.__version__)

# Prefetch用にどの程度先取りしておくか自動調整

#　これもいつもどおりでOK

# https://tensorflow.classcat.com/2019/03/23/tf20-alpha-guide-data-performance/

AUTO = tf.data.experimental.AUTOTUNE

# Kaggleのデータ拝借

from kaggle_datasets import KaggleDatasets

TPUとGPUを切り替える部分。
各TPUを制御するcluster_resolverを使ってTPUをネットワークから探索。
各TPU間の同期方法を定義するstrategyは必須。strategyはあくまでインスタンスなので、のちほど定義する
TPUStrategyの場合、アクセラレーターで分散する際に変数のミラーリングがされ、分散学習の重みを集約→重みを更新(同期)→共有する。この仕組みをall-reduceっていう。
TPUは使えない場合はGPUの同じミラーリング戦略であるMirroredStrategyをつかう
以下、日経 xTECHを拝借

try:
    tpu = tf.distribute.cluster_resolver.TPUClusterResolver.connect() # TPU detection
    strategy = tf.distribute.TPUStrategy(tpu) 
    print("Detect: TPU")
except ValueError:
     # detect GPUs
    strategy = tf.distribute.MirroredStrategy() # for GPU or multi-GPU machines
    #strategy = tf.distribute.get_strategy() # default strategy that works on CPU and single GPU(インストールで入れた際に)
    #strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy() # for clusters of multi-GPU machines
    print("Detect: GPU")

print("Number of accelerators: ", strategy.num_replicas_in_sync)
# tf.distribute.Strategyは、トレーニングを複数のGPU、複数のマシン、またはTPUに分散するためのTensorFlowAPI
# https://tensorflow.classcat.com/2019/03/21/tf20-alpha-guide-distribute-strategy/

try:

tpu = tf.distribute.cluster_resolver.TPUClusterResolver.connect() # TPU detection

strategy = tf.distribute.TPUStrategy(tpu)

print("Detect: TPU")

except ValueError:

# detect GPUs

strategy = tf.distribute.MirroredStrategy() # for GPU or multi-GPU machines

#strategy = tf.distribute.get_strategy() # default strategy that works on CPU and single GPU(インストールで入れた際に)

#strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy() # for clusters of multi-GPU machines

print("Detect: GPU")

print("Number of accelerators: ", strategy.num_replicas_in_sync)

# tf.distribute.Strategyは、トレーニングを複数のGPU、複数のマシン、またはTPUに分散するためのTensorFlowAPI

# https://tensorflow.classcat.com/2019/03/21/tf20-alpha-guide-distribute-strategy/

まずはデータを撮ってくる。
バッチサイズが大事で、ここは８の倍数(GoogleTPUv3.8の場合8bit量子化しているので)にしないと行けない。
これは後で使うSteps_per_executionも同じことが言える。

"""
Define image info
"""
# Dataset: https://www.kaggle.com/mgornergoogle/five-flowers
# TFRecord ファイル群
EPOCHS = 12
IMAGE_SIZE = [331, 331]

# available image sizes
# train/test/val
FLOWERS_DATASETS = { 
    192: GCS_PATH + '/tfrecords-jpeg-192x192/*.tfrec',
    224: GCS_PATH + '/tfrecords-jpeg-224x224/*.tfrec',
    331: GCS_PATH + '/tfrecords-jpeg-331x331/*.tfrec',
    512: GCS_PATH + '/tfrecords-jpeg-512x512/*.tfrec'
}
# Labels
CLASSES = ['daisy', 'dandelion', 'roses', 'sunflowers', 'tulips']
assert IMAGE_SIZE[0] == IMAGE_SIZE[1], "only square images are supported"
assert IMAGE_SIZE[0] in FLOWERS_DATASETS, "this image size is not supported"

# strategy.num_replicas_in_sync: #TPU  or #GPU: レプリカを各GPU/TPUで作成
BATCH_SIZE = 16 * strategy.num_replicas_in_sync

"""

Define image info

"""

# Dataset: https://www.kaggle.com/mgornergoogle/five-flowers

# TFRecord ファイル群

EPOCHS = 12

IMAGE_SIZE = [331, 331]

# available image sizes

# train/test/val

FLOWERS_DATASETS = {

192: GCS_PATH + '/tfrecords-jpeg-192x192/*.tfrec',

224: GCS_PATH + '/tfrecords-jpeg-224x224/*.tfrec',

331: GCS_PATH + '/tfrecords-jpeg-331x331/*.tfrec',

512: GCS_PATH + '/tfrecords-jpeg-512x512/*.tfrec'

}

# Labels

CLASSES = ['daisy', 'dandelion', 'roses', 'sunflowers', 'tulips']

assert IMAGE_SIZE[0] == IMAGE_SIZE[1], "only square images are supported"

assert IMAGE_SIZE[0] in FLOWERS_DATASETS, "this image size is not supported"

# strategy.num_replicas_in_sync: #TPU or #GPU: レプリカを各GPU/TPUで作成

BATCH_SIZE = 16 * strategy.num_replicas_in_sync

学習率を途中で変えるらしいです。TPUに限定してなぜカスタマイズするのかは不明。
ただ自分自身で変更するのは初めて。よい知見。

"""
learning rate scheduler
"""

LR_START = 0.00001
LR_MAX = 0.00005 * strategy.num_replicas_in_sync
LR_MIN = 0.00001
LR_RAMPUP_EPOCHS = 5
LR_SUSTAIN_EPOCHS = 0
LR_EXP_DECAY = .8

def lrfn(epoch):
    if epoch < LR_RAMPUP_EPOCHS:
        lr = (LR_MAX - LR_START) / LR_RAMPUP_EPOCHS * epoch + LR_START
    elif epoch < LR_RAMPUP_EPOCHS + LR_SUSTAIN_EPOCHS:
        lr = LR_MAX
    else:
        #　以降指数的に減衰
        lr = (LR_MAX - LR_MIN) * LR_EXP_DECAY**(epoch - LR_RAMPUP_EPOCHS - LR_SUSTAIN_EPOCHS) + LR_MIN
    return lr
    
lr_callback = tf.keras.callbacks.LearningRateScheduler(lrfn, verbose=True)

rng = [i for i in range(EPOCHS)]
y = [lrfn(x) for x in rng]
plt.plot(rng, y)
print("Learning rate schedule: {:.3g} to {:.3g} to {:.3g}".format(y[0], max(y), y[-1]))

"""

learning rate scheduler

"""

LR_START = 0.00001

LR_MAX = 0.00005 * strategy.num_replicas_in_sync

LR_MIN = 0.00001

LR_RAMPUP_EPOCHS = 5

LR_SUSTAIN_EPOCHS = 0

LR_EXP_DECAY = .8

def lrfn(epoch):

if epoch < LR_RAMPUP_EPOCHS:

lr = (LR_MAX - LR_START) / LR_RAMPUP_EPOCHS * epoch + LR_START

elif epoch < LR_RAMPUP_EPOCHS + LR_SUSTAIN_EPOCHS:

lr = LR_MAX

else:

#　以降指数的に減衰

lr = (LR_MAX - LR_MIN) * LR_EXP_DECAY**(epoch - LR_RAMPUP_EPOCHS - LR_SUSTAIN_EPOCHS) + LR_MIN

return lr

lr_callback = tf.keras.callbacks.LearningRateScheduler(lrfn, verbose=True)

rng = [i for i in range(EPOCHS)]

y = [lrfn(x) for x in rng]

plt.plot(rng, y)

print("Learning rate schedule: {:.3g} to {:.3g} to {:.3g}".format(y[0], max(y), y[-1]))

この辺は画像を撮ってくるだけ。TFRecordなので、それに応じて読み込む。
ここは全て借り物。おパクり。
重要な点は最初に設定した”AUTO”。
TPU・GPUはPrefetchが速度のボトルネックなのでここは大事。それ以外はデータよりけりだと思う

"""
Data Loading Funcs
全て借り物関数
"""
#TPUはサイズに気をつけないといけないらしい
#https://cloud.google.com/tpu/docs/performance-guide#tf-fcts
# https://qiita.com/ohtaman/items/325615358f008a4005e2

def count_data_items(filenames):
    # the number of data items is written in the name of the .tfrec files, i.e. flowers00-230.tfrec = 230 data items
    n = [int(re.compile(r"-([0-9]*)\.").search(filename).group(1)) for filename in filenames]
    return np.sum(n)
        
def read_tfrecord(example):
    features = {
        "image": tf.io.FixedLenFeature([], tf.string), # tf.string means bytestring
        "class": tf.io.FixedLenFeature([], tf.int64),  # shape [] means scalar
        "one_hot_class": tf.io.VarLenFeature(tf.float32),# Load variable length feature
    }
    example = tf.io.parse_single_example(example, features)
    image = tf.image.decode_jpeg(example['image'], channels=3) # pixel format uint8 [0,255] range
    class_label = tf.cast(example['class'], tf.int32) # not used
    one_hot_class = tf.sparse.to_dense(example['one_hot_class'])
    one_hot_class = tf.reshape(one_hot_class, [5])
    return image, one_hot_class


def force_image_sizes(dataset, image_size):
    # explicit size needed for TPU
    reshape_images = lambda image, label: (tf.reshape(image, [*image_size, 3]), label)
    dataset = dataset.map(reshape_images, num_parallel_calls=AUTO)
    return dataset

def load_dataset(filenames):
    # Read from TFRecords. For optimal performance, reading from multiple files at once and
    # disregarding data order. Order does not matter since we will be shuffling the data anyway.

    ignore_order = tf.data.Options()
    ignore_order.experimental_deterministic = False

    dataset = tf.data.TFRecordDataset(filenames, num_parallel_reads=AUTO) # automatically interleaves reads from multiple files
    dataset = dataset.with_options(ignore_order) # uses data as soon as it streams in, rather than in its original order
    dataset = dataset.map(read_tfrecord, num_parallel_calls=AUTO)
    dataset = force_image_sizes(dataset, IMAGE_SIZE)
    return dataset

def data_augment(image, one_hot_class):
    # data augmentation. Thanks to the dataset.prefetch(AUTO) statement in the next function (below),
    # this happens essentially for free on TPU. Data pipeline code is executed on the "CPU" part
    # of the TPU while the TPU itself is computing gradients.
    image = tf.image.random_flip_left_right(image)
    image = tf.image.random_saturation(image, 0, 2)
    return image, one_hot_class   

def get_training_dataset():
    dataset = load_dataset(TRAINING_FILENAMES)
    dataset = dataset.map(data_augment, num_parallel_calls=AUTO)
    dataset = dataset.repeat()
    dataset = dataset.shuffle(2048)
    dataset = dataset.batch(BATCH_SIZE)
    dataset = dataset.prefetch(AUTO) # prefetch next batch while training (autotune prefetch buffer size)
    return dataset

def get_validation_dataset():
    dataset = load_dataset(VALIDATION_FILENAMES)
    dataset = dataset.batch(BATCH_SIZE)
    dataset = dataset.prefetch(AUTO) # prefetch next batch while training (autotune prefetch buffer size)
    return dataset

"""

Data Loading Funcs

全て借り物関数

"""

#TPUはサイズに気をつけないといけないらしい

#https://cloud.google.com/tpu/docs/performance-guide#tf-fcts

# https://qiita.com/ohtaman/items/325615358f008a4005e2

def count_data_items(filenames):

# the number of data items is written in the name of the .tfrec files, i.e. flowers00-230.tfrec = 230 data items

n = [int(re.compile(r"-([0-9]*)\.").search(filename).group(1)) for filename in filenames]

return np.sum(n)

def read_tfrecord(example):

features = {

"image": tf.io.FixedLenFeature([], tf.string), # tf.string means bytestring

"class": tf.io.FixedLenFeature([], tf.int64), # shape [] means scalar

"one_hot_class": tf.io.VarLenFeature(tf.float32),# Load variable length feature

}

example = tf.io.parse_single_example(example, features)

image = tf.image.decode_jpeg(example['image'], channels=3) # pixel format uint8 [0,255] range

class_label = tf.cast(example['class'], tf.int32) # not used

one_hot_class = tf.sparse.to_dense(example['one_hot_class'])

one_hot_class = tf.reshape(one_hot_class, [5])

return image, one_hot_class

def force_image_sizes(dataset, image_size):

# explicit size needed for TPU

reshape_images = lambda image, label: (tf.reshape(image, [*image_size, 3]), label)

dataset = dataset.map(reshape_images, num_parallel_calls=AUTO)

return dataset

def load_dataset(filenames):

# Read from TFRecords. For optimal performance, reading from multiple files at once and

# disregarding data order. Order does not matter since we will be shuffling the data anyway.

ignore_order = tf.data.Options()

ignore_order.experimental_deterministic = False

dataset = tf.data.TFRecordDataset(filenames, num_parallel_reads=AUTO) # automatically interleaves reads from multiple files

dataset = dataset.with_options(ignore_order) # uses data as soon as it streams in, rather than in its original order

dataset = dataset.map(read_tfrecord, num_parallel_calls=AUTO)

dataset = force_image_sizes(dataset, IMAGE_SIZE)

return dataset

def data_augment(image, one_hot_class):

# data augmentation. Thanks to the dataset.prefetch(AUTO) statement in the next function (below),

# this happens essentially for free on TPU. Data pipeline code is executed on the "CPU" part

# of the TPU while the TPU itself is computing gradients.

image = tf.image.random_flip_left_right(image)

image = tf.image.random_saturation(image, 0, 2)

return image, one_hot_class

def get_training_dataset():

dataset = load_dataset(TRAINING_FILENAMES)

dataset = dataset.map(data_augment, num_parallel_calls=AUTO)

dataset = dataset.repeat()

dataset = dataset.shuffle(2048)

dataset = dataset.batch(BATCH_SIZE)

dataset = dataset.prefetch(AUTO) # prefetch next batch while training (autotune prefetch buffer size)

return dataset

def get_validation_dataset():

dataset = load_dataset(VALIDATION_FILENAMES)

dataset = dataset.batch(BATCH_SIZE)

dataset = dataset.prefetch(AUTO) # prefetch next batch while training (autotune prefetch buffer size)

return dataset

よしなに分割。

"""
Split Data
"""
gcs_pattern = FLOWERS_DATASETS[IMAGE_SIZE[0]]
validation_split = 0.19
filenames = tf.io.gfile.glob(gcs_pattern)
split = len(filenames) - int(len(filenames) * validation_split)
TRAINING_FILENAMES = filenames[:split]
VALIDATION_FILENAMES = filenames[split:]
TRAIN_STEPS = count_data_items(TRAINING_FILENAMES) // BATCH_SIZE
VALIDATION_STEPS = -(-count_data_items(VALIDATION_FILENAMES) // BATCH_SIZE) # The "-(-//)" trick rounds up instead of down :-)
print("TRAINING IMAGES: ", count_data_items(TRAINING_FILENAMES), ", STEPS PER EPOCH: ", TRAIN_STEPS)
print("VALIDATION IMAGES: ", count_data_items(VALIDATION_FILENAMES))

"""

Split Data

"""

gcs_pattern = FLOWERS_DATASETS[IMAGE_SIZE[0]]

validation_split = 0.19

filenames = tf.io.gfile.glob(gcs_pattern)

split = len(filenames) - int(len(filenames) * validation_split)

TRAINING_FILENAMES = filenames[:split]

VALIDATION_FILENAMES = filenames[split:]

TRAIN_STEPS = count_data_items(TRAINING_FILENAMES) // BATCH_SIZE

VALIDATION_STEPS = -(-count_data_items(VALIDATION_FILENAMES) // BATCH_SIZE) # The "-(-//)" trick rounds up instead of down :-)

print("TRAINING IMAGES: ", count_data_items(TRAINING_FILENAMES), ", STEPS PER EPOCH: ", TRAIN_STEPS)

print("VALIDATION IMAGES: ", count_data_items(VALIDATION_FILENAMES))

モデル作成。

strategy.scope()でくくることで、先程のインスタンスがようやく登場。
この中の処理のみがTPUで共有される。つまり、モデル部分だけが共有される。
今回の例ではSAMPLEに従って、xceptionを採用。下記の通り他のモデルを使いたいならテキトウに変更
steps_per_execution=8は、1callでなんバッチ読み込むかを定義する。上の”AUTO”と同じで高速化に貢献
単純なギモンとしてGPUでも同じこと言えるよね？と思ったが、GPUは関係ないらしい。

# xceptionを利用
# https://note.nkmk.me/python-tensorflow-keras-applications-pretrained-models/
#　VGGとかこの部分を変えればよい。
# ベつに自作してもよい。普通の画像分類
with strategy.scope():
    """
    Distributed Execution
    """
    # xception用のｎ前処理実行
    img_adjust_layer = tf.keras.layers.Lambda(
        lambda data: tf.keras.applications.xception.preprocess_input(
            tf.cast(data, tf.float32)), input_shape=[*IMAGE_SIZE, 3])
    # xceptionモデルの組み込み
    pretrained_model = tf.keras.applications.Xception(weights='imagenet', include_top=False)
    #学習モード
    pretrained_model.trainable = True

    model = tf.keras.Sequential([
        img_adjust_layer,
        pretrained_model,
        #  GAP>Flat>Denceで終端処理
        tf.keras.layers.GlobalAveragePooling2D(),
        tf.keras.layers.Flatten(),
        tf.keras.layers.Dense(5, activation='softmax')
    ])

# steps_per_execution: 1回のrunで何バッチを呼ぶか？
# TPUまたは小さいモデルなら効果あり（50%高速化）。
# GPUがなぜ効果内のか不明。
#https://keras.io/api/models/model_training_apis/
# https://qiita.com/T-STAR/items/e2998d4c22c882039ffb
model.compile(
    optimizer='adam',
    loss = 'categorical_crossentropy',
    metrics=['accuracy'],
　　steps_per_execution=8
)

model.summary()

# xceptionを利用

# https://note.nkmk.me/python-tensorflow-keras-applications-pretrained-models/

#　VGGとかこの部分を変えればよい。

# ベつに自作してもよい。普通の画像分類

with strategy.scope():

"""

Distributed Execution

"""

# xception用のｎ前処理実行

img_adjust_layer = tf.keras.layers.Lambda(

lambda data: tf.keras.applications.xception.preprocess_input(

tf.cast(data, tf.float32)), input_shape=[*IMAGE_SIZE, 3])

# xceptionモデルの組み込み

pretrained_model = tf.keras.applications.Xception(weights='imagenet', include_top=False)

#学習モード

pretrained_model.trainable = True

model = tf.keras.Sequential([

img_adjust_layer,

pretrained_model,

# GAP>Flat>Denceで終端処理

tf.keras.layers.GlobalAveragePooling2D(),

tf.keras.layers.Flatten(),

tf.keras.layers.Dense(5, activation='softmax')

])

# steps_per_execution: 1回のrunで何バッチを呼ぶか？

# TPUまたは小さいモデルなら効果あり（50%高速化）。

# GPUがなぜ効果内のか不明。

#https://keras.io/api/models/model_training_apis/

# https://qiita.com/T-STAR/items/e2998d4c22c882039ffb

model.compile(

optimizer='adam',

loss = 'categorical_crossentropy',

metrics=['accuracy'],

　　steps_per_execution=8

)

model.summary()

学習。もはやGPUのときと変わらない。

history = model.fit(get_training_dataset(), steps_per_epoch=TRAIN_STEPS, epochs=EPOCHS,
                    validation_data=get_validation_dataset(), validation_steps=VALIDATION_STEPS,
                    callbacks=[lr_callback])

final_accuracy = history.history["val_accuracy"][-5:]
print("FINAL ACCURACY MEAN-5: ", np.mean(final_accuracy))

history = model.fit(get_training_dataset(), steps_per_epoch=TRAIN_STEPS, epochs=EPOCHS,

validation_data=get_validation_dataset(), validation_steps=VALIDATION_STEPS,

callbacks=[lr_callback])

final_accuracy = history.history["val_accuracy"][-5:]

print("FINAL ACCURACY MEAN-5: ", np.mean(final_accuracy))

このあと保存と読み込みがあるが、GPUと読み込み方法は異なるものの、余り本質ではないんので省略。
いまはTenswflow liteというものがあって、いわゆるエッジコンピューティング用に特化されたフレームワークがあるけど、そこでも活用できるみたい。
電子工作でおなじみのswitchScienceから出ている。今度遊んでみたい。1万するけどｗｗｗ
https://www.switch-science.com/catalog/5817/

gmoがなんかやっていた。

https://recruit.gmo.jp/engineer/jisedai/blog/edge_tpu_tf_lite_basics/