机器学习笔记 - 学习使用TensorFlow和张量处理单元 (TPU) 构建图像分类模型

1、概述

这篇文章是基于kaggle官方的教程简单整理而来，目的是为了人了解并使用张量处理单元 (TPU) 创建基线模型。数据集基于木薯叶病分类竞赛的数据。我们将使用 TensorFlow 和 Keras 来构建我们的计算机视觉模型，并使用 TPU 来训练我们的模型并进行预测。

原文地址：Getting Started: TPUs + Cassava Leaf Disease | Kaggle

张量处理单元 (TPU)

张量处理单元 (TPU) 是专门用于深度学习任务的硬件加速器，是谷歌专门为神经网络机器学习开发的人工智能加速器专用集成电路(ASIC) ，特别是使用谷歌自己的TensorFlow软件。所有 Kaggler 每周都有 30 小时的免费 TPU 时间，并且可以在单次会话中使用最多 3 小时。

TPU 硬件

大约 20 英寸（50 厘米）的 TPU v3-8 板是相当大的硬件。它采用 4 个双核 TPU 芯片，总共 8 个 TPU 内核。

每个 TPU 内核都有一个传统的矢量处理部分 (VPU) 以及能够处理 128x128 矩阵的专用矩阵乘法硬件。这是专门加速机器学习工作负载的部分。

TPU 配备 128GB 的高速内存，允许更大的批次、更大的模型以及更大的训练输入。在上面的示例中，您可以尝试使用数据集中也提供的 512x512 px 输入图像，并查看 TPU v3-8 轻松处理它们。

2、设置环境/检测TPU

打印tensorflow版本，代码用到了kaggle的库，需要在kaggle提供的环境上运行。

import math, re, os
import tensorflow as tf
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from kaggle_datasets import KaggleDatasets
from tensorflow import keras
from functools import partial
from sklearn.model_selection import train_test_split
print("Tensorflow version " + tf.__version__)

我们在这里使用我们的代码所做的是确保我们将通过 TPU 发送我们的数据。您要查找的是副本数：8 的打印输出，对应于 TPU 的 8 个核心。如果您的打印输出显示副本数：1，则您的笔记本中可能没有启用 TPU。

try:tpu = tf.distribute.cluster_resolver.TPUClusterResolver()print('Device:', tpu.master())tf.config.experimental_connect_to_cluster(tpu)tf.tpu.experimental.initialize_tpu_system(tpu)strategy = tf.distribute.experimental.TPUStrategy(tpu)
except:strategy = tf.distribute.get_strategy()
print('Number of replicas:', strategy.num_replicas_in_sync)

如果显示类似下面，则是TPU

Device: grpc://10.0.0.2:8470
Number of replicas: 8

3、设置相关变量

AUTOTUNE = tf.data.experimental.AUTOTUNE
GCS_PATH = KaggleDatasets().get_gcs_path()
BATCH_SIZE = 16 * strategy.num_replicas_in_sync
IMAGE_SIZE = [512, 512]
CLASSES = ['0', '1', '2', '3', '4']
EPOCHS = 25

我们正在处理的数据已被格式化为 TFRecords，这是一种用于存储二进制记录序列的格式。 TFRecords 与 TPU 配合得非常好，允许我们通过 TPU 发送少量大文件进行处理。

因为我们的数据仅包含训练和测试图像，所以下面我们将使用 train_test_split() 函数将训练数据拆分为训练和验证数据。

4、解码数据

在下面的代码块中，我们将设置一系列函数，允许我们将图像转换为张量，以便我们可以在模型中使用它们。我们还将规范化我们的数据。我们的图像使用范围为 [0, 255] 的“红、蓝、绿 (RBG)”比例，通过对其进行归一化，我们将每个像素的值设置为 [0, 1] 范围内的数字 .

def decode_image(image):image = tf.image.decode_jpeg(image, channels=3)image = tf.cast(image, tf.float32) / 255.0image = tf.reshape(image, [*IMAGE_SIZE, 3])return image

如果您回想机器学习简介，您可能还记得我们是如何设置变量 X 和 y 来表示我们的特征 X 和预测目标 y。这段代码完成了类似的事情，尽管我们的特征不是使用标签 X 和 y，而是由术语图像表示，我们的预测目标由术语目标表示。

您可能还会注意到，此函数考虑了未标记的图像。这是因为我们的测试图像没有任何标签。

def read_tfrecord(example, labeled):tfrecord_format = {"image": tf.io.FixedLenFeature([], tf.string),"target": tf.io.FixedLenFeature([], tf.int64)} if labeled else {"image": tf.io.FixedLenFeature([], tf.string),"image_name": tf.io.FixedLenFeature([], tf.string)}example = tf.io.parse_single_example(example, tfrecord_format)image = decode_image(example['image'])if labeled:label = tf.cast(example['target'], tf.int32)return image, labelidnum = example['image_name']return image, idnum

我们将使用以下函数来加载我们的数据集。 TPU 的优点之一是我们可以一次跨 TPU 运行多个文件，这说明了使用 TPU 的速度优势。为了利用这一点，我们希望确保在数据流入后立即使用数据，而不是造成数据流瓶颈。

def load_dataset(filenames, labeled=True, ordered=False):ignore_order = tf.data.Options()if not ordered:ignore_order.experimental_deterministic = False # disable order, increase speeddataset = tf.data.TFRecordDataset(filenames, num_parallel_reads=AUTOTUNE) # automatically interleaves reads from multiple filesdataset = dataset.with_options(ignore_order) # uses data as soon as it streams in, rather than in its original orderdataset = dataset.map(partial(read_tfrecord, labeled=labeled), num_parallel_calls=AUTOTUNE)return dataset

虽然使用 train_test_split() 创建训练和验证数据集，但请考虑探索交叉验证。

TRAINING_FILENAMES, VALID_FILENAMES = train_test_split(tf.io.gfile.glob(GCS_PATH + '/train_tfrecords/ld_train*.tfrec'),test_size=0.35, random_state=5
)TEST_FILENAMES = tf.io.gfile.glob(GCS_PATH + '/test_tfrecords/ld_test*.tfrec')

5、增强数据

您可以在 TensorFlow tf.image 文档中阅读有关这些增强（以及所有其他可用的增强！）的更多信息。

def data_augment(image, label):# Thanks to the dataset.prefetch(AUTO) statement in the following function 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)return image, label

6、定义数据加载方法

以下函数将用于加载我们的训练、验证和测试数据集，并打印出每个数据集中的图像数量。

def get_training_dataset():dataset = load_dataset(TRAINING_FILENAMES, labeled=True)  dataset = dataset.map(data_augment, num_parallel_calls=AUTOTUNE)  dataset = dataset.repeat()dataset = dataset.shuffle(2048)dataset = dataset.batch(BATCH_SIZE)dataset = dataset.prefetch(AUTOTUNE)return datasetdef get_validation_dataset(ordered=False):dataset = load_dataset(VALID_FILENAMES, labeled=True, ordered=ordered) dataset = dataset.batch(BATCH_SIZE)dataset = dataset.cache()dataset = dataset.prefetch(AUTOTUNE)return datasetdef get_test_dataset(ordered=False):dataset = load_dataset(TEST_FILENAMES, labeled=False, ordered=ordered)dataset = dataset.batch(BATCH_SIZE)dataset = dataset.prefetch(AUTOTUNE)return datasetdef count_data_items(filenames):n = [int(re.compile(r"-([0-9]*)\.").search(filename).group(1)) for filename in filenames]return np.sum(n)NUM_TRAINING_IMAGES = count_data_items(TRAINING_FILENAMES)
NUM_VALIDATION_IMAGES = count_data_items(VALID_FILENAMES)
NUM_TEST_IMAGES = count_data_items(TEST_FILENAMES)print('Dataset: {} training images, {} validation images, {} (unlabeled) test images'.format(NUM_TRAINING_IMAGES, NUM_VALIDATION_IMAGES, NUM_TEST_IMAGES))

打印如下

Dataset: 13380 training images, 8017 validation images, 1 (unlabeled) test images

7、构建模型

学习率调度

lr_scheduler = keras.optimizers.schedules.ExponentialDecay(initial_learning_rate=1e-5, decay_steps=10000, decay_rate=0.9)

为了确保我们的模型在TPU上进行训练，我们使用 strategy.scope() 构建它。

该模型是使用迁移学习构建的，这意味着我们有一个预训练模型 (ResNet50) 作为我们的基础模型，然后是使用 tf.keras.Sequential 构建的可定制模型。如果您是迁移学习的新手，我建议将 base_model.trainable 设置为 False，但鼓励您更改正在使用的基本模型（更多选项可在 tf.keras.applications 模块文档中找到）并在定制模型。

请注意，我们使用 sparse_categorical_crossentropy 作为损失函数，因为我们没有对标签进行 one-hot 编码。

with strategy.scope():       img_adjust_layer = tf.keras.layers.Lambda(tf.keras.applications.resnet50.preprocess_input, input_shape=[*IMAGE_SIZE, 3])base_model = tf.keras.applications.ResNet50(weights='imagenet', include_top=False)base_model.trainable = Falsemodel = tf.keras.Sequential([tf.keras.layers.BatchNormalization(renorm=True),img_adjust_layer,base_model,tf.keras.layers.GlobalAveragePooling2D(),tf.keras.layers.Dense(8, activation='relu'),#tf.keras.layers.BatchNormalization(renorm=True),tf.keras.layers.Dense(len(CLASSES), activation='softmax')  ])model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=lr_scheduler, epsilon=0.001),loss='sparse_categorical_crossentropy',  metrics=['sparse_categorical_accuracy'])

8、训练模型

当我们的模型正在训练时，您会看到每个 epoch 的打印输出，还可以通过单击笔记本右上角工具栏中的 TPU 指标来监控 TPU 使用情况。

# load data
train_dataset = get_training_dataset()
valid_dataset = get_validation_dataset()

EPS_PER_EPOCH = NUM_TRAINING_IMAGES // BATCH_SIZE
VALID_STEPS = NUM_VALIDATION_IMAGES // BATCH_SIZEhistory = model.fit(train_dataset, steps_per_epoch=STEPS_PER_EPOCH, epochs=EPOCHS,validation_data=valid_dataset,validation_steps=VALID_STEPS)