语义分割与 SegFormer 和 Hugging Face Transformers
作者: Sayak Paul
创建日期: 2023/01/25
最后修改: 2023/01/29
描述: 微调一个 SegFormer 模型变体用于语义分割。
简介
在这个例子中，我们展示如何微调一个 SegFormer 模型变体以在自定义数据集上进行语义分割。语义分割是将每个像素分配一个类别的任务。SegFormer 在 SegFormer: Simple and Efficient Design for Semantic Segmentation with Transformers 中提出。SegFormer 使用层次化的 Transformer 架构（称为 "Mix Transformer"）作为其编码器，以及一个轻量级的解码器进行分割。因此，它在语义分割上提供了最先进的性能，同时比现有模型更高效。欲了解更多细节，请参阅原始论文。
我们利用 Hugging Face Transformers 加载预训练的 SegFormer 检查点，并在自定义数据集上进行微调。
注意： 此示例重用了以下来源的代码：
要运行此示例，我们需要安装 transformers 库：
!!pip install transformers -q
[]
加载数据
我们使用 Oxford-IIIT Pets 数据集作为这个例子的基础。我们利用 tensorflow_datasets 加载该数据集。
import tensorflow_datasets as tfds
dataset, info = tfds.load("oxford_iiit_pet:3.*.*", with_info=True)
/opt/conda/lib/python3.7/site-packages/tensorflow_io/python/ops/__init__.py:98: UserWarning: unable to load libtensorflow_io_plugins.so: unable to open file: libtensorflow_io_plugins.so, from paths: ['/opt/conda/lib/python3.7/site-packages/tensorflow_io/python/ops/libtensorflow_io_plugins.so']
caused by: ['/opt/conda/lib/python3.7/site-packages/tensorflow_io/python/ops/libtensorflow_io_plugins.so: undefined symbol: _ZN3tsl5mutexC1Ev']
warnings.warn(f"unable to load libtensorflow_io_plugins.so: {e}")
/opt/conda/lib/python3.7/site-packages/tensorflow_io/python/ops/__init__.py:104: UserWarning: file system plugins are not loaded: unable to open file: libtensorflow_io.so, from paths: ['/opt/conda/lib/python3.7/site-packages/tensorflow_io/python/ops/libtensorflow_io.so']
caused by: ['/opt/conda/lib/python3.7/site-packages/tensorflow_io/python/ops/libtensorflow_io.so: undefined symbol: _ZNK10tensorflow4data11DatasetBase8FinalizeEPNS_15OpKernelContextESt8functionIFN3tsl8StatusOrISt10unique_ptrIS1_NS5_4core15RefCountDeleterEEEEvEE']
warnings.warn(f"file system plugins are not loaded: {e}")
准备数据集
为准备训练和评估所需的数据集，我们：
- 使用在预训练 SegFormer 中使用的均值和标准差来标准化图像。
- 从分割掩码中减去 1，以使像素值从 0 开始。
- 调整图像大小。
- 转置图像，使其为
"channels_first"格式。这是为了使其与 Hugging Face Transformers 的 SegFormer 模型兼容。
import tensorflow as tf
from tensorflow.keras import backend
image_size = 512
mean = tf.constant([0.485, 0.456, 0.406])
std = tf.constant([0.229, 0.224, 0.225])
def normalize(input_image, input_mask):
input_image = tf.image.convert_image_dtype(input_image, tf.float32)
input_image = (input_image - mean) / tf.maximum(std, backend.epsilon())
input_mask -= 1
return input_image, input_mask
def load_image(datapoint):
input_image = tf.image.resize(datapoint["image"], (image_size, image_size))
input_mask = tf.image.resize(
datapoint["segmentation_mask"],
(image_size, image_size),
method="bilinear",
)
input_image, input_mask = normalize(input_image, input_mask)
input_image = tf.transpose(input_image, (2, 0, 1))
return {"pixel_values": input_image, "labels": tf.squeeze(input_mask)}
我们现在使用上述工具准备 tf.data.Dataset 对象，包括
prefetch() 提高性能。更改 batch_size 以匹配您用于训练的 GPU 的内存大小。
auto = tf.data.AUTOTUNE
batch_size = 4
train_ds = (
dataset["train"]
.cache()
.shuffle(batch_size * 10)
.map(load_image, num_parallel_calls=auto)
.batch(batch_size)
.prefetch(auto)
)
test_ds = (
dataset["test"]
.map(load_image, num_parallel_calls=auto)
.batch(batch_size)
.prefetch(auto)
)
我们可以检查输入图像及其分割图的形状：
print(train_ds.element_spec)
{'pixel_values': TensorSpec(shape=(None, 3, 512, 512), dtype=tf.float32, name=None), 'labels': TensorSpec(shape=(None, 512, 512), dtype=tf.float32, name=None)}
可视化数据集
import matplotlib.pyplot as plt
def display(display_list):
plt.figure(figsize=(15, 15))
title = ["输入图像", "真实掩膜", "预测掩膜"]
for i in range(len(display_list)):
plt.subplot(1, len(display_list), i + 1)
plt.title(title[i])
plt.imshow(tf.keras.utils.array_to_img(display_list[i]))
plt.axis("off")
plt.show()
for samples in train_ds.take(2):
sample_image, sample_mask = samples["pixel_values"][0], samples["labels"][0]
sample_image = tf.transpose(sample_image, (1, 2, 0))
sample_mask = tf.expand_dims(sample_mask, -1)
display([sample_image, sample_mask])
加载预训练的 SegFormer 检查点
我们现在从 Hugging Face Transformers 加载一个预训练的 SegFormer 模型变体。SegFormer 模型有多个不同的变体，称为 MiT-B0 到 MiT-B5。您可以在 这里 找到这些检查点。 我们加载最小的变体 Mix-B0，它在推理效率和预测性能之间提供了良好的折中。
from transformers import TFSegformerForSemanticSegmentation
model_checkpoint = "nvidia/mit-b0"
id2label = {0: "外部", 1: "内部", 2: "边界"}
label2id = {label: id for id, label in id2label.items()}
num_labels = len(id2label)
model = TFSegformerForSemanticSegmentation.from_pretrained(
model_checkpoint,
num_labels=num_labels,
id2label=id2label,
label2id=label2id,
ignore_mismatched_sizes=True,
)
WARNING:tensorflow:最近的 5 次调用
WARNING:tensorflow:最近的 5 次调用
WARNING:tensorflow:最近的 6 次调用
WARNING:tensorflow:最近的 6 次调用
TFSegformerForSemanticSegmentation
正在初始化一个新的解码器头。
我们现在可以初始化一个优化器并用它来编译模型。
编译模型
lr = 0.00006
optimizer = tf.keras.optimizers.Adam(learning_rate=lr)
model.compile(optimizer=optimizer)
编译时未指定损失 - 模型的内部损失计算将被用作损失。不要惊慌 - 这是在 Transformers 中训练 TensorFlow 模型的常见方法！要禁用此行为，请传递损失参数，或者如果您不希望模型计算损失，则明确传递 `loss=None`。
请注意，我们在编译模型时没有使用任何损失函数。这是因为模型的前向传递
实现
了当我们提供标签和输入图像时的损失计算部分。计算损失后，模型返回一个结构化的 dataclass 对象，然后用来指导训练过程。
随着编译好的模型，我们可以继续并调用 fit() 来开始微调过程！
预测回调以监控训练进度
这帮助我们在模型被微调时可视化一些样本预测，从而帮助我们监控模型的进展。此回调的灵感来自 这个教程。
from IPython.display import clear_output
def create_mask(pred_mask):
pred_mask = tf.math.argmax(pred_mask, axis=1)
pred_mask = tf.expand_dims(pred_mask, -1)
return pred_mask[0]
def show_predictions(dataset=None, num=1):
if dataset:
for sample in dataset.take(num):
images, masks = sample["pixel_values"], sample["labels"]
masks = tf.expand_dims(masks, -1)
pred_masks = model.predict(images).logits
images = tf.transpose(images, (0, 2, 3, 1))
display([images[0], masks[0], create_mask(pred_masks)])
else:
display(
[
sample_image,
sample_mask,
create_mask(model.predict(tf.expand_dims(sample_image, 0))),
]
)
class DisplayCallback(tf.keras.callbacks.Callback):
def __init__(self, dataset, **kwargs):
super().__init__(**kwargs)
self.dataset = dataset
def on_epoch_end(self, epoch, logs=None):
clear_output(wait=True)
show_predictions(self.dataset)
print("\n在第 {} 轮结束后的样本预测\n".format(epoch + 1))
训练模型
# 如果结果没有达到预期质量，请增加训练轮数。
epochs = 5
history = model.fit(
train_ds,
validation_data=test_ds,
callbacks=[DisplayCallback(test_ds)],
epochs=epochs,
)
1/1 [==============================] - 0s 54ms/step
第 5 轮结束后的样本预测
920/920 [==============================] - 89s 97ms/step - loss: 0.1742 - val_loss: 0.1927
推理
我们对测试集中的一些样本进行推理。
show_predictions(test_ds, 5)
1/1 [==============================] - 0s 54ms/step
1/1 [==============================] - 0s 54ms/step
1/1 [==============================] - 0s 53ms/step
1/1 [==============================] - 0s 53ms/step
1/1 [==============================] - 0s 53ms/step
结论
在这个例子中，我们学习了如何在一个自定义数据集上微调 SegFormer 模型变体以进行语义分割。为了简洁，示例保持简短。然而，还有几个事情，你可以进一步尝试：
* 结合数据增强以潜在地提高结果。
* 使用更大的SegFormer模型检查点，以观察结果如何受到影响。
* 将微调后的模型推送到Hugging Face，以便轻松与社区分享。
你只需执行 model.push_to_hub("your-username/your-awesome-model") 即可。
然后你可以通过以下方式加载模型
TFSegformerForSemanticSegmentation.from_pretrained("your-username/your-awesome-model")。
这里
是一个端到端的示例，如果你在寻找参考。
* 如果你更愿意在模型微调的同时将模型检查点推送到Hub，你可以使用 PushToHubCallback Keras回调。
这里是一个示例。
这里是使用此回调创建的模型
库的示例。