使用 TensorFlow Similarity 进行图像相似性搜索的度量学习

pip -q install tensorflow_similarity

import random

from matplotlib import pyplot as plt
from mpl_toolkits import axes_grid1
import numpy as np

import tensorflow as tf
from tensorflow import keras

import tensorflow_similarity as tfsim


tfsim.utils.tf_cap_memory()

print("TensorFlow:", tf.__version__)
print("TensorFlow Similarity:", tfsim.__version__)

# 这决定了训练期间使用的类别数量。
# 在这里我们使用所有类别。
num_known_classes = 10
class_list = random.sample(population=range(10), k=num_known_classes)

classes_per_batch = 10
# 每个批次传递每个类别的多个示例确保每个示例都有
# 多个正对。进行三元组挖掘或
# 使用诸如`MultiSimilarityLoss`或`CircleLoss`这类损失时，这可能是有用的，因为这些可以
# 采用所有正对的加权混合。一般来说，每个
# 类别的更多示例将为正对提供更多信息，而每个批次的更多类别
# 将在负对中提供更多不同的信息。然而，
# 损失计算批次中示例之间的成对距离，因此
# 批次大小的上限受限于内存。
examples_per_class_per_batch = 8

print(
    "批次大小为: "
    f"{min(classes_per_batch, num_known_classes) * examples_per_class_per_batch}"
)

print(" 创建训练数据 ".center(34, "#"))
train_ds = tfsim.samplers.TFDatasetMultiShotMemorySampler(
    "cifar10",
    classes_per_batch=min(classes_per_batch, num_known_classes),
    splits="train",
    steps_per_epoch=4000,
    examples_per_class_per_batch=examples_per_class_per_batch,
    class_list=class_list,
)

print("\n" + " 创建验证数据 ".center(34, "#"))
val_ds = tfsim.samplers.TFDatasetMultiShotMemorySampler(
    "cifar10",
    classes_per_batch=classes_per_batch,
    splits="test",
    total_examples_per_class=100,
)

num_cols = num_rows = 5
# 获取前25个示例。
x_slice, y_slice = train_ds.get_slice(begin=0, size=num_cols * num_rows)

fig = plt.figure(figsize=(6.0, 6.0))
grid = axes_grid1.ImageGrid(fig, 111, nrows_ncols=(num_cols, num_rows), axes_pad=0.1)

for ax, im, label in zip(grid, x_slice, y_slice):
    ax.imshow(im)
    ax.axis("off")

embedding_size = 256

inputs = keras.layers.Input((32, 32, 3))
x = keras.layers.Rescaling(scale=1.0 / 255)(inputs)
x = keras.layers.Conv2D(64, 3, activation="relu")(x)
x = keras.layers.BatchNormalization()(x)
x = keras.layers.Conv2D(128, 3, activation="relu")(x)
x = keras.layers.BatchNormalization()(x)
x = keras.layers.MaxPool2D((4, 4))(x)
x = keras.layers.Conv2D(256, 3, activation="relu")(x)
x = keras.layers.BatchNormalization()(x)
x = keras.layers.Conv2D(256, 3, activation="relu")(x)
x = keras.layers.GlobalMaxPool2D()(x)
outputs = tfsim.layers.MetricEmbedding(embedding_size)(x)

# 构建模型
model = tfsim.models.SimilarityModel(inputs, outputs)
model.summary()

epochs = 3
learning_rate = 0.002
val_steps = 50

# 初始化相似度损失
loss = tfsim.losses.MultiSimilarityLoss()

# 编译和训练
model.compile(
    optimizer=keras.optimizers.Adam(learning_rate), loss=loss, steps_per_execution=10,
)
history = model.fit(
    train_ds, epochs=epochs, validation_data=val_ds, validation_steps=val_steps
)

x_index, y_index = val_ds.get_slice(begin=0, size=200)
model.reset_index()
model.index(x_index, y_index, data=x_index)

x_train, y_train = train_ds.get_slice(begin=0, size=1000)
calibration = model.calibrate(
    x_train,
    y_train,
    calibration_metric="f1",
    matcher="match_nearest",
    extra_metrics=["precision", "recall", "binary_accuracy"],
    verbose=1,
)

num_neighbors = 5
labels = [
    "飞机",
    "汽车",
    "鸟",
    "猫",
    "鹿",
    "狗",
    "青蛙",
    "马",
    "船",
    "卡车",
    "未知",
]
class_mapping = {c_id: c_lbl for c_id, c_lbl in zip(range(11), labels)}

x_display, y_display = val_ds.get_slice(begin=200, size=10)
# 在索引中查找最近的邻居
nns = model.lookup(x_display, k=num_neighbors)

# 显示
for idx in np.argsort(y_display):
    tfsim.visualization.viz_neigbors_imgs(
        x_display[idx],
        y_display[idx],
        nns[idx],
        class_mapping=class_mapping,
        fig_size=(16, 2),
    )

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 5))
x = calibration.thresholds["distance"]

ax1.plot(x, calibration.thresholds["precision"], label="精确度")
ax1.plot(x, calibration.thresholds["recall"], label="召回率")
ax1.plot(x, calibration.thresholds["f1"], label="F1分数")
ax1.legend()
ax1.set_title("距离增加时指标的演变")
ax1.set_xlabel("距离")
ax1.set_ylim((-0.05, 1.05))

ax2.plot(calibration.thresholds["recall"], calibration.thresholds["precision"])
ax2.set_title("精确度-召回率曲线")
ax2.set_xlabel("召回率")
ax2.set_ylabel("精确度")
ax2.set_ylim((-0.05, 1.05))
plt.show()

cutpoint = "optimal"

# 这为每个类别提供100个示例。
# 我们在创建val_ds采样器时定义了这个。
x_confusion, y_confusion = val_ds.get_slice(0, -1)

matches = model.match(x_confusion, cutpoint=cutpoint, no_match_label=10)
cm = tfsim.visualization.confusion_matrix(
    matches,
    y_confusion,
    labels=labels,
    title="切点:%s的混淆矩阵" % cutpoint,
    normalize=False,
)

idx_no_match = np.where(np.array(matches) == 10)
no_match_queries = x_confusion[idx_no_match]
if len(no_match_queries):
    plt.imshow(no_match_queries[0])
else:
    print("所有查询在距离阈值以下都有匹配。")

# val_ds 中每个类别被限制为 100 个示例。
num_examples_to_clusters = 1000
thumb_size = 96
plot_size = 800
vx, vy = val_ds.get_slice(0, num_examples_to_clusters)

# 取消注释以运行交互式投影仪。
# tfsim.visualization.projector(
#     model.predict(vx),
#     labels=vy,
#     images=vx,
#     class_mapping=class_mapping,
#     image_size=thumb_size,
#     plot_size=plot_size,
# )