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

1、概述

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

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

        张量处理单元 (TPU)

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

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.0
    image = 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, label
    idnum = 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 speed
    dataset = tf.data.TFRecordDataset(filenames, num_parallel_reads=AUTOTUNE) # 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(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 dataset

def 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 dataset

def get_test_dataset(ordered=False):
    dataset = load_dataset(TEST_FILENAMES, labeled=False, ordered=ordered)
    dataset = dataset.batch(BATCH_SIZE)
    dataset = dataset.prefetch(AUTOTUNE)
    return dataset

def 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 = False
    
    model = 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_SIZE

history = model.fit(train_dataset, 
                    steps_per_epoch=STEPS_PER_EPOCH, 
                    epochs=EPOCHS,
                    validation_data=valid_dataset,
                    validation_steps=VALID_STEPS)