引言

本文介绍了一种双编码器（也称作双塔）神经网络模型的应用，该模型利用自然语言处理技术实现图像搜索。该模型受到Alec Radford等人提出的CLIP方法的启发，其核心思想是结合训练一个视觉编码器和一个文本编码器，将图像和其标题的表示投射到同一个嵌入空间，使得标题嵌入接近其描述的图像嵌入。

为了运行此模型，需要安装一些必要的库，包括TensorFlow 2.4及以上版本，TensorFlow Hub，TensorFlow Text，以及TensorFlow Addons。这些库可以通过以下命令安装：

bash pip install -q -U tensorflow-hub tensorflow-text tensorflow-addons

数据准备

本例中使用的数据集是MS-COCO，该数据集包含了超过82,000张图片，每张图片至少有5个不同的标题注释。我们重新应用这些图像标题对来训练双编码器模型，以实现图像搜索功能。

数据下载和提取

首先，需要下载并解压数据集，该数据集由两个压缩文件夹组成：一个是图像，另一个是对应的图像标题。压缩后的图像文件夹大小约为13GB。以下是数据下载和解压的代码示例：

```python rootdir = "datasets" annotationsdir = os.path.join(rootdir, "annotations") imagesdir = os.path.join(rootdir, "train2014") tfrecordsdir = os.path.join(rootdir, "tfrecords") annotationfile = os.path.join(annotationsdir, "captionstrain2014.json")

下载并解压标注文件

if not os.path.exists(annotationsdir): annotationzip = tf.keras.utils.getfile( "captions.zip", cachedir=os.path.abspath("."), origin="http://images.cocodataset.org/annotations/annotationstrainval2014.zip", extract=True, ) os.remove(annotationzip)

下载并解压图像文件

if not os.path.exists(imagesdir): imagezip = tf.keras.utils.getfile( "train2014.zip", cachedir=os.path.abspath("."), origin="http://images.cocodataset.org/zips/train2014.zip", extract=True, ) os.remove(image_zip)

print("数据集已成功下载并解压。") with open(annotationfile, "r") as f: annotations = json.load(f)["annotations"] imagepathtocaption = collections.defaultdict(list) for element in annotations: caption = f"{element['caption'].lower().rstrip('.')}" imagepath = imagesdir + "/COCOtrain2014" + "%012d.jpg" % (element["imageid"]) imagepathtocaption[imagepath].append(caption) imagepaths = list(imagepathtocaption.keys()) print(f"图像数量：{len(imagepaths)}") ```

数据处理及存储

接下来，我们处理并保存数据到TFRecord文件中。在此过程中，可以选择不同的样本量来控制用于训练双编码器模型的图像-标题对的数量。在这个例子中，我们将训练集设为30,000张图像，验证集设为5,000张图像，每张图像使用两个标题，从而产生60,000个图像-标题对。

```python trainsize = 30000 validsize = 5000 captionsperimage = 2 imagesperfile = 2000 trainimagepaths = imagepaths[:trainsize] numtrainfiles = int(np.ceil(trainsize / imagesperfile)) trainfilesprefix = os.path.join(tfrecordsdir, "train") validimagepaths = imagepaths[-validsize:] numvalidfiles = int(np.ceil(validsize / imagesperfile)) validfilesprefix = os.path.join(tfrecordsdir, "valid")

tf.io.gfile.makedirs(tfrecords_dir)

def bytesfeature(value): return tf.train.Feature(byteslist=tf.train.BytesList(value=[value]))

def createexample(imagepath, caption): feature = { "caption": bytesfeature(caption.encode()), "rawimage": bytesfeature(tf.io.readfile(image_path).numpy()), } return tf.train.Example(features=tf.train.Features(feature=feature))

def writetfrecords(filename, imagepaths): captionlist = [] imagepathlist = [] for imagepath in imagepaths: captions = imagepathtocaption[imagepath][:captionsperimage] captionlist.extend(captions) imagepathlist.extend([imagepath] * len(captions)) with tf.io.TFRecordWriter(filename) as writer: for exampleidx in range(len(imagepathlist)): example = createexample( imagepathlist[exampleidx], captionlist[exampleidx] ) writer.write(example.SerializeToString()) return example_idx + 1

def writedata(imagepaths, numfiles, filesprefix): examplecounter = 0 for fileidx in tqdm(range(numfiles)): filename = filesprefix + "-%02d.tfrecord" % (fileidx) startidx = imagesperfile * fileidx endidx = startidx + imagesperfile examplecounter += writetfrecords(filename, imagepaths[startidx:endidx]) return example_counter

trainexamplecount = writedata(trainimagepaths, numtrainfiles, trainfilesprefix) print(f"{trainexample_count} 训练样本已写入TFRecord文件。")

validexamplecount = writedata(validimagepaths, numvalidfiles, validfilesprefix) print(f"{validexample_count} 验证样本已写入TFRecord文件。") ```

创建数据集

我们从TFRecord文件中读取数据并创建tf.data.Dataset对象。该数据集用于训练和评估双编码器模型。

```python featuredescription = { "caption": tf.io.FixedLenFeature([], tf.string), "rawimage": tf.io.FixedLenFeature([], tf.string), }

def readexample(example): features = tf.io.parsesingleexample(example, featuredescription) rawimage = features.pop("rawimage") features["image"] = tf.image.resize( tf.image.decodejpeg(rawimage, channels=3), size=(299, 299) ) return features

def getdataset(filepattern, batchsize): return ( tf.data.TFRecordDataset(tf.data.Dataset.listfiles(filepattern)) .map( readexample, numparallelcalls=tf.data.experimental.AUTOTUNE, deterministic=False, ) .shuffle(batchsize * 10) .prefetch(buffersize=tf.data.experimental.AUTOTUNE) .batch(batch_size) ) ```

实时投影头

投影头用于将图像和文本嵌入到具有相同维度的嵌入空间中。下面是定义投影头的方法：

python def project_embeddings( embeddings, num_projection_layers, projection_dims, dropout_rate ): projected_embeddings = layers.Dense(units=projection_dims)(embeddings) for _ in range(num_projection_layers): x = tf.nn.gelu(projected_embeddings) x = layers.Dense(projection_dims)(x) x = layers.Dropout(dropout_rate)(x) x = layers.Add()([projected_embeddings, x]) projected_embeddings = layers.LayerNormalization()(x) return projected_embeddings

完成视觉编码器

在本例中，我们使用Keras Applications中的Xception作为视觉编码器的基础。以下是创建视觉编码器的方法：

python def create_vision_encoder( num_projection_layers, projection_dims, dropout_rate, trainable=False ): xception = keras.applications.Xception( include_top=False, weights="imagenet", pooling="avg" ) for layer in xception.layers: layer.trainable = trainable inputs = layers.Input(shape=(299, 299, 3), name="image_input") xception_input = tf.keras.applications.xception.preprocess_input(inputs) embeddings = xception(xception_input) outputs = project_embeddings( embeddings, num_projection_layers, projection_dims, dropout_rate ) return keras.Model(inputs, outputs, name="vision_encoder")

完成文本编码器

我们使用TensorFlow Hub中的BERT作为文本编码器。以下是创建文本编码器的方法：

python def create_text_encoder( num_projection_layers, projection_dims, dropout_rate, trainable=False ): preprocess = hub.KerasLayer( "https://tfhub.dev/tensorflow/bert_en_uncased_preprocess/2", name="text_preprocessing" ) bert = hub.KerasLayer( "https://tfhub.dev/tensorflow/small_bert/bert_en_uncased_L-4_H-512_A-8/1", "bert" ) bert.trainable = trainable inputs = layers.Input(shape=(), dtype=tf.string, name="text_input") bert_inputs = preprocess(inputs) embeddings = bert(bert_inputs)["pooled_output"] outputs = project_embeddings( embeddings, num_projection_layers, projection_dims, dropout_rate ) return keras.Model(inputs, outputs, name="text_encoder")

完成双编码器

为了计算损失，我们计算每个captioni和imagej之间的对偶点积相似度作为预测值。captioni和imagej之间的目标相似度计算为（captioni和captionj之间的点积相似度）和（imagei和imagej之间的点积相似度）的平均值。然后，我们使用交叉熵来计算目标和预测之间的损失。

```python class DualEncoder(keras.Model): def init(self, textencoder, imageencoder, temperature=1.0, kwargs): super(DualEncoder, self)._init(kwargs) self.textencoder = textencoder self.imageencoder = imageencoder self.temperature = temperature self.losstracker = keras.metrics.Mean(name="loss")

@property
def metrics(self):
    return [self.loss_tracker]

def call(self, features, training=False):
    with tf.device("/gpu:0"):
        caption_embeddings = self.text_encoder(features["caption"], training=training)
    with tf.device("/gpu:1"):
        image_embeddings = self.image_encoder(features["image"], training=training)
    return caption_embeddings, image_embeddings

def compute_loss(self, caption_embeddings, image_embeddings):
    logits = (
        tf.matmul(caption_embeddings, image_embeddings, transpose_b=True) / self.temperature
    )
    images_similarity = tf.matmul(
        image_embeddings, image_embeddings, transpose_b=True
    )
    captions_similarity = tf.matmul(
        caption_embeddings, caption_embeddings, transpose_b=True
    )
    targets = keras.activations.softmax(
        (captions_similarity + images_similarity) / (2 * self.temperature)
    )
    captions_loss = keras.losses.categorical_crossentropy(
        y_true=targets, y_pred=logits, from_logits=True
    )
    images_loss = keras.losses.categorical_crossentropy(
        y_true=tf.transpose(targets), y_pred=tf.transpose(logits), from_logits=True
    )
    return (captions_loss + images_loss) / 2

def train_step(self, features):
    with tf.GradientTape() as tape:
        caption_embeddings, image_embeddings = self(features, training=True)
        loss = self.compute_loss(caption_embeddings, image_embeddings)
    gradients = tape.gradient(loss, self.trainable_variables)
    self.optimizer.apply_gradients(zip(gradients, self.trainable_variables))
    self.loss_tracker.update_state(loss)
    return {"loss": self.loss_tracker.result()}

def test_step(self, features):
    caption_embeddings, image_embeddings = self(features, training=False)
    loss = self.compute_loss(caption_embeddings, image_embeddings)
    self.loss_tracker.update_state(loss)
    return {"loss": self.loss_tracker.result()}

```

训练双编码器模型

在这个实验中，我们冻结了文本和图像的基础编码器，只让投影头停止训练。

```python numepochs = 5 # 实际中应训练至少30个周期 batchsize = 256 visionencoder = createvisionencoder( numprojectionlayers=1, projectiondims=256, dropoutrate=0.1 ) textencoder = createtextencoder( numprojectionlayers=1, projectiondims=256, dropoutrate=0.1 ) dualencoder = DualEncoder(textencoder, visionencoder, temperature=0.05) dualencoder.compile( optimizer=tfa.optimizers.AdamW(learningrate=0.001, weightdecay=0.001) )

print(f"GPU数量：{len(tf.config.listphysicaldevices('GPU'))}") print(f"样本数量（caption-image pairs）：{trainexamplecount}") print(f"批量大小：{batchsize}") print(f"每轮迭代次数：{int(np.ceil(trainexamplecount / batchsize))}")

traindataset = getdataset(os.path.join(tfrecordsdir, "train-*.tfrecord"), batchsize) validdataset = getdataset(os.path.join(tfrecordsdir, "valid-*.tfrecord"), batchsize)

reducelr = keras.callbacks.ReduceLROnPlateau( monitor="valloss", factor=0.2, patience=3 ) earlystopping = tf.keras.callbacks.EarlyStopping( monitor="valloss", patience=5, restorebestweights=True )

history = dualencoder.fit( traindataset, epochs=numepochs, validationdata=validdataset, callbacks=[reducelr, earlystopping], ) print("训练完成。保存视觉和文本编码器...") visionencoder.save("visionencoder") textencoder.save("text_encoder") print("模型已保存。") ```

使用自然语言查询搜索图像

训练完成后，我们可以通过以下步骤来检索对应自然语言查询的图像：

1. 将图像输入到视觉编码器中，生成图像的嵌入。
2. 将自然语言查询输入到文本编码器中，生成查询嵌入。
3. 计算查询嵌入与索引中的图像嵌入之间的相似度，以检索出最匹配的图像。
4. 查看前k个匹配图像的路径，并将其显示出来。

```python print("加载视觉和文本编码器...") visionencoder = keras.models.loadmodel("visionencoder") textencoder = keras.models.loadmodel("textencoder") print("模型已加载。")

def readimage(imagepath): imagearray = tf.image.decodejpeg(tf.io.readfile(imagepath), channels=3) return tf.image.resize(image_array, (299, 299))

print(f"生成{len(imagepaths)}张图像的嵌入...") imageembeddings = visionencoder.predict( tf.data.Dataset.fromtensorslices(imagepaths).map(readimage).batch(batchsize), verbose=1, ) print(f"图像嵌入形状：{image_embeddings.shape}。")

def findmatches(imageembeddings, queries, k=9, normalize=True): queryembedding = textencoder(tf.converttotensor(queries)) if normalize: imageembeddings = tf.math.l2normalize(imageembeddings, axis=1) queryembedding = tf.math.l2normalize(queryembedding, axis=1) dotsimilarity = tf.matmul(queryembedding, imageembeddings, transposeb=True) results = tf.math.topk(dotsimilarity, k).indices.numpy() return [[image_paths[idx] for idx in indices] for indices in results]

query = "a family standing next to the ocean on a sandy beach with a surf board" matches = findmatches(imageembeddings, [query], normalize=True)[0] plt.figure(figsize=(20, 20)) for i in range(9): ax = plt.subplot(3, 3, i + 1) plt.imshow(mpimg.imread(matches[i])) plt.axis("off") ```

评价检索质量

为了评价双编码器模型，我们使用标题作为查询。使用训练外样本图像和标题来评价检索质量，使用top k精度。如果对于一个给定的标题，其相关的图像在前k个匹配范围内被检索到，则算作一个正确的预测。

```python def computetopkaccuracy(imagepaths, k=100): hits = 0 numbatches = int(np.ceil(len(imagepaths) / batchsize)) for idx in tqdm(range(numbatches)): startidx = idx * batchsize endidx = startidx + batchsize currentimagepaths = imagepaths[startidx:endidx] queries = [ imagepathtocaption[imagepath][0] for imagepath in currentimagepaths ] result = findmatches(imageembeddings, queries, k) hits += sum( [ imagepath in matches for (imagepath, matches) in list(zip(currentimagepaths, result)) ] ) return hits / len(imagepaths)

print("评分训练数据...") trainaccuracy = computetopkaccuracy(trainimagepaths) print(f"训练准确率：{round(train_accuracy * 100, 3)}%")

print("评分验证数据...") evalaccuracy = computetopkaccuracy(imagepaths[trainsize:]) print(f"验证准确率：{round(eval_accuracy * 100, 3)}%") ```

结论

通过增加训练样本数量、训练更多周期、探索其他图像和文本的基础编码器、设置基础编码器的可训练性，以及调整超参数，特别是softmax温度损失计算，可以获得更好的结果。