|
2021年1月、OpenAIはDALL-EとCLIPという2つの新しいモデルを発表しました。どちらも、何らかの方法でテキストと画像を結び付けるマルチモーダルモデルです。 CLIP は Contrastive Language–Image Pre-training の略で、対照的なテキストと画像のペアに基づく事前トレーニング方法です。 CLIPを導入する理由人気のある安定拡散は単一のモデルではなく、複数のモデルで構成されているためです。テキスト エンコーダーは、ユーザーのテキスト入力をエンコードするために使用されます。このテキスト エンコーダーは、CLIP モデルのテキスト エンコーダーです。 CLIP モデルをトレーニングするときに、入力文を与えると、それに合わせて最も関連性の高い画像が抽出されます。 CLIP は完全な文とそれが説明する画像との関係を学習します。つまり、「車」、「犬」などの個別のカテゴリではなく、完全な文章に基づいてトレーニングされており、これはアプリケーションにとって非常に重要です。完全なフレーズでトレーニングすると、モデルはより多くのことを学習し、写真とテキスト間のパターンを認識できるようになります。また、かなり大規模な写真とそれに対応する文章のデータセットでトレーニングすると、モデルが分類器として機能できることも実証しました。 CLIP がリリースされたとき、ImageNet データセットで微調整を行った後、微調整なし (ゼロショット) で ResNets-50 を上回るパフォーマンスを発揮することができ、非常に有用であることが分かりました。 そこでこの記事では、CLIP をより深く理解できるように、PyTorch で CLIP モデルを最初から実装します。 ここでは、timm と transformers という 2 つのライブラリを使用する必要があります。まずはコードをインポートしましょう。 import os import cv2 import gc import numpy as np import pandas as pd import itertools from tqdm.autonotebook import tqdm import albumentations as A import matplotlib.pyplot as plt import torch from torch import nn import torch.nn.functional as F import timm from transformers import DistilBertModel, DistilBertConfig, DistilBertTokenizer 次のステップは、データと一般的な構成を前処理することです。 config は、すべてのハイパーパラメータを配置する通常の Python ファイルです。Jupyter Notebook の場合は、ノートブックの先頭で定義されたクラスです。 class CFG: debug = False image_path = "../input/flickr-image-dataset/flickr30k_images/flickr30k_images" captions_path = "." batch_size = 32 num_workers = 4 head_lr = 1e-3 image_encoder_lr = 1e-4 text_encoder_lr = 1e-5 weight_decay = 1e-3 patience = 1 factor = 0.8 epochs = 2 device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model_name = 'resnet50' image_embedding = 2048 text_encoder_model = "distilbert-base-uncased" text_embedding = 768 text_tokenizer = "distilbert-base-uncased" max_length = 200 pretrained = True # for both image encoder and text encoder trainable = True # for both image encoder and text encoder temperature = 1.0 # image size size = 224 # for projection head; used for both image and text encoders num_projection_layers = 1 projection_dim = 256 dropout = 0.1 カスタムインジケーター用の補助クラスもいくつかあります class AvgMeter: def __init__(self, name="Metric"): self.name = name self.reset() def reset(self): self.avg, self.sum, self.count = [0] * 3 def update(self, val, count=1): self.count += count self.sum += val * count self.avg = self.sum / self.count def __repr__(self): text = f"{self.name}: {self.avg:.4f}" return text def get_lr(optimizer): for param_group in optimizer.param_groups: return param_group["lr"]
私たちの目標は画像と文章を説明することです。したがって、データセットは文章と画像の両方を返す必要があります。したがって、DistilBERT トークナイザーを使用して文 (タイトル) をトークン化し、トークン ID (input_ids) とアテンション マスクを DistilBERT に提供する必要があります。 DistilBERT は BERT モデルよりも小さいですが、モデルの結果は似ているため、DistilBERT を使用することを選択します。 次のステップは、HuggingFace トークナイザーを使用してトークン化することです。 __init__ で取得されたトークナイザー オブジェクトは、モデルの実行時にロードされます。ヘッダーは、事前に決められた最大長にパディングされ、切り捨てられます。関連する画像を読み込む前に、エンコードされたキャプションを__getitem__に読み込みます。これは、キー input_ids と attention_mask を持つ辞書であり、それらを変換して拡張します (存在する場合)。次に、これはテンソルに変換され、「画像」をキーとして辞書に保存されます。最後に、タイトルの生のテキストをキーワード「title」とともに辞書に入力します。 class CLIPDataset(torch.utils.data.Dataset): def __init__(self, image_filenames, captions, tokenizer, transforms): """ image_filenames and cpations must have the same length; so, if there are multiple captions for each image, the image_filenames must have repetitive file names """ self.image_filenames = image_filenames self.captions = list(captions) self.encoded_captions = tokenizer( list(captions), padding=True, truncatinotallow=True, max_length=CFG.max_length ) self.transforms = transforms def __getitem__(self, idx): item = { key: torch.tensor(values[idx]) for key, values in self.encoded_captions.items() } image = cv2.imread(f"{CFG.image_path}/{self.image_filenames[idx]}") image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) image = self.transforms(image=image)['image'] item['image'] = torch.tensor(image).permute(2, 0, 1).float() item['caption'] = self.captions[idx] return item def __len__(self): return len(self.captions) def get_transforms(mode="train"): if mode == "train": return A.Compose( [ A.Resize(CFG.size, CFG.size, always_apply=True), A.Normalize(max_pixel_value=255.0, always_apply=True), ] ) else: return A.Compose( [ A.Resize(CFG.size, CFG.size, always_apply=True), A.Normalize(max_pixel_value=255.0, always_apply=True), ] ) 画像およびテキスト エンコーダー: 画像エンコーダーとして ResNet50 を使用します。 class ImageEncoder(nn.Module): """ Encode images to a fixed size vector """ def __init__( self, model_name=CFG.model_name, pretrained=CFG.pretrained, trainable=CFG.trainable ): super().__init__() self.model = timm.create_model( model_name, pretrained, num_classes=0, global_pool="avg" ) for p in self.model.parameters(): p.requires_grad = trainable def forward(self, x): return self.model(x) テキスト エンコーダーとして DistilBERT を使用します。 CLS トークンの最終的な表現は、文の全体の表現を取得するために使用されます。 class TextEncoder(nn.Module): def __init__(self, model_name=CFG.text_encoder_model, pretrained=CFG.pretrained, trainable=CFG.trainable): super().__init__() if pretrained: self.model = DistilBertModel.from_pretrained(model_name) else: self.model = DistilBertModel(cnotallow=DistilBertConfig()) for p in self.model.parameters(): p.requires_grad = trainable # we are using the CLS token hidden representation as the sentence's embedding self.target_token_idx = 0 def forward(self, input_ids, attention_mask): output = self.model(input_ids=input_ids, attention_mask=attention_mask) last_hidden_state = output.last_hidden_state return last_hidden_state[:, self.target_token_idx, :] 上記のコードでは、画像とテキストを固定サイズのベクトル (画像の場合は 2048、テキストの場合は 768) にエンコードしています。画像とテキストを比較するには、同じ次元にする必要があります。そのため、2048 次元と 768 次元のベクトルを 256 次元 (projection_dim) に投影します。次元が同じ場合にのみ比較できます。 class ProjectionHead(nn.Module): def __init__( self, embedding_dim, projection_dim=CFG.projection_dim, dropout=CFG.dropout ): super().__init__() self.projection = nn.Linear(embedding_dim, projection_dim) self.gelu = nn.GELU() self.fc = nn.Linear(projection_dim, projection_dim) self.dropout = nn.Dropout(dropout) self.layer_norm = nn.LayerNorm(projection_dim) def forward(self, x): projected = self.projection(x) x = self.gelu(projected) x = self.fc(x) x = self.dropout(x) x = x + projected x = self.layer_norm(x) return x 最終的に、CLIP モデルは次のようになります。 class CLIPModel(nn.Module): def __init__( self, temperature=CFG.temperature, image_embedding=CFG.image_embedding, text_embedding=CFG.text_embedding, ): super().__init__() self.image_encoder = ImageEncoder() self.text_encoder = TextEncoder() self.image_projection = ProjectionHead(embedding_dim=image_embedding) self.text_projection = ProjectionHead(embedding_dim=text_embedding) self.temperature = temperature def forward(self, batch): # Getting Image and Text Features image_features = self.image_encoder(batch["image"]) text_features = self.text_encoder( input_ids=batch["input_ids"], attention_mask=batch["attention_mask"] ) # Getting Image and Text Embeddings (with same dimension) image_embeddings = self.image_projection(image_features) text_embeddings = self.text_projection(text_features) # Calculating the Loss logits = (text_embeddings @ image_embeddings.T) / self.temperature images_similarity = image_embeddings @ image_embeddings.T texts_similarity = text_embeddings @ text_embeddings.T targets = F.softmax( (images_similarity + texts_similarity) / 2 * self.temperature, dim=-1 ) texts_loss = cross_entropy(logits, targets, reductinotallow='none') images_loss = cross_entropy(logits.T, targets.T, reductinotallow='none') loss = (images_loss + texts_loss) / 2.0 # shape: (batch_size) return loss.mean() #这里还加了一个交叉熵函数def cross_entropy(preds, targets, reductinotallow='none'): log_softmax = nn.LogSoftmax(dim=-1) loss = (-targets * log_softmax(preds)).sum(1) if reduction == "none": return loss elif reduction == "mean": return loss.mean() ここで注目すべきは、CLIP は対称クロスエントロピーを損失関数として使用しており、これによりノイズの影響が軽減され、モデルの堅牢性が向上するということです。簡単にするために、ここではクロスエントロピーのみを使用します。 テストできる内容: # A simple Example batch_size = 4 dim = 256 embeddings = torch.randn(batch_size, dim) out = embeddings @ embeddings.T print(F.softmax(out, dim=-1)) 次のステップはトレーニングです。トレーニングと検証のためにデータローダーをロードするのに役立つ関数がいくつかあります。 def make_train_valid_dfs(): dataframe = pd.read_csv(f"{CFG.captions_path}/captions.csv") max_id = dataframe["id"].max() + 1 if not CFG.debug else 100 image_ids = np.arange(0, max_id) np.random.seed(42) valid_ids = np.random.choice( image_ids, size=int(0.2 * len(image_ids)), replace=False ) train_ids = [id_ for id_ in image_ids if id_ not in valid_ids] train_dataframe = dataframe[dataframe["id"].isin(train_ids)].reset_index(drop=True) valid_dataframe = dataframe[dataframe["id"].isin(valid_ids)].reset_index(drop=True) return train_dataframe, valid_dataframe def build_loaders(dataframe, tokenizer, mode): transforms = get_transforms(mode=mode) dataset = CLIPDataset( dataframe["image"].values, dataframe["caption"].values, tokenizer=tokenizer, transforms=transforms, ) dataloader = torch.utils.data.DataLoader( dataset, batch_size=CFG.batch_size, num_workers=CFG.num_workers, shuffle=True if mode == "train" else False, ) return dataloader そしてトレーニングと評価があります def train_epoch(model, train_loader, optimizer, lr_scheduler, step): loss_meter = AvgMeter() tqdm_object = tqdm(train_loader, total=len(train_loader)) for batch in tqdm_object: batch = {k: v.to(CFG.device) for k, v in batch.items() if k != "caption"} loss = model(batch) optimizer.zero_grad() loss.backward() optimizer.step() if step == "batch": lr_scheduler.step() count = batch["image"].size(0) loss_meter.update(loss.item(), count) tqdm_object.set_postfix(train_loss=loss_meter.avg, lr=get_lr(optimizer)) return loss_meter def valid_epoch(model, valid_loader): loss_meter = AvgMeter() tqdm_object = tqdm(valid_loader, total=len(valid_loader)) for batch in tqdm_object: batch = {k: v.to(CFG.device) for k, v in batch.items() if k != "caption"} loss = model(batch) count = batch["image"].size(0) loss_meter.update(loss.item(), count) tqdm_object.set_postfix(valid_loss=loss_meter.avg) return loss_meter
最後に、プロセス全体が統合されます def main(): train_df, valid_df = make_train_valid_dfs() tokenizer = DistilBertTokenizer.from_pretrained(CFG.text_tokenizer) train_loader = build_loaders(train_df, tokenizer, mode="train") valid_loader = build_loaders(valid_df, tokenizer, mode="valid") model = CLIPModel().to(CFG.device) params = [ {"params": model.image_encoder.parameters(), "lr": CFG.image_encoder_lr}, {"params": model.text_encoder.parameters(), "lr": CFG.text_encoder_lr}, {"params": itertools.chain( model.image_projection.parameters(), model.text_projection.parameters() ), "lr": CFG.head_lr, "weight_decay": CFG.weight_decay} ] optimizer = torch.optim.AdamW(params, weight_decay=0.) lr_scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau( optimizer, mode="min", patience=CFG.patience, factor=CFG.factor ) step = "epoch" best_loss = float('inf') for epoch in range(CFG.epochs): print(f"Epoch: {epoch + 1}") model.train() train_loss = train_epoch(model, train_loader, optimizer, lr_scheduler, step) model.eval() with torch.no_grad(): valid_loss = valid_epoch(model, valid_loader) if valid_loss.avg < best_loss: best_loss = valid_loss.avg torch.save(model.state_dict(), "best.pt") print("Saved Best Model!") lr_scheduler.step(valid_loss.avg)
アプリケーション: 画像の埋め込みを取得し、一致するものを見つけます。 トレーニングを完了した後、実際にどのように適用するのでしょうか?トレーニング済みのモデルを読み込み、検証セットから画像を供給し、形状 (valid_set_size, 256) の image_embeddings とモデル自体を返す関数を記述する必要があります。 def get_image_embeddings(valid_df, model_path): tokenizer = DistilBertTokenizer.from_pretrained(CFG.text_tokenizer) valid_loader = build_loaders(valid_df, tokenizer, mode="valid") model = CLIPModel().to(CFG.device) model.load_state_dict(torch.load(model_path, map_locatinotallow=CFG.device)) model.eval() valid_image_embeddings = [] with torch.no_grad(): for batch in tqdm(valid_loader): image_features = model.image_encoder(batch["image"].to(CFG.device)) image_embeddings = model.image_projection(image_features) valid_image_embeddings.append(image_embeddings) return model, torch.cat(valid_image_embeddings) _, valid_df = make_train_valid_dfs() model, image_embeddings = get_image_embeddings(valid_df, "best.pt") def find_matches(model, image_embeddings, query, image_filenames, n=9): tokenizer = DistilBertTokenizer.from_pretrained(CFG.text_tokenizer) encoded_query = tokenizer([query]) batch = { key: torch.tensor(values).to(CFG.device) for key, values in encoded_query.items() } with torch.no_grad(): text_features = model.text_encoder( input_ids=batch["input_ids"], attention_mask=batch["attention_mask"] ) text_embeddings = model.text_projection(text_features) image_embeddings_n = F.normalize(image_embeddings, p=2, dim=-1) text_embeddings_n = F.normalize(text_embeddings, p=2, dim=-1) dot_similarity = text_embeddings_n @ image_embeddings_n.T values, indices = torch.topk(dot_similarity.squeeze(0), n * 5) matches = [image_filenames[idx] for idx in indices[::5]] _, axes = plt.subplots(3, 3, figsize=(10, 10)) for match, ax in zip(matches, axes.flatten()): image = cv2.imread(f"{CFG.image_path}/{match}") image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) ax.imshow(image) ax.axis("off") plt.show() 呼び出し方法は次のとおりです。 find_matches(model, image_embeddings, query="one dog sitting on the grass", image_filenames=valid_df['image'].values, n=9) カスタマイズ効果がかなり良いことがわかります（ただし、写真には猫がいます、笑）。つまり、CLIP メソッドは小さなデータ セットでもカスタマイズできます。 この記事のコードとデータセットは次のとおりです。 https://www.kaggle.com/code/jyotidabas/simple-openai-clip-implementation |
