目次
前に書いてある
私は自己教師あり学習を始めたばかりで、自己教師あり学習についての理解はまだ理論段階ですが、この自己教師あり学習コードリーディング集という自分自身に風穴を開けたいと思っています。同時に、私のような初心者にも役立つようになりたいと思っていますので、何か問題があれば、皆さんにアドバイスをいただければ幸いです。
1.mocoのメインアイデア
moco について話す前に、対照学習とは何かを知る必要があります。対照学習は自己教師あり学習の重要な分野です。自己教師あり学習は、独自のデータセットで教師あり情報をマイニングし、それ自体で生成された教師あり情報を通じてモデルをトレーニングします。たとえば、画像を 9 つのグリッドに切り取り、ラベルを付けます各グリッドに 1 ～ 9 を割り当て、9 つのグリッドとラベルをスクランブルし、スクランブルされた画像を入力として、スクランブルされたラベルをグラウンド トゥルースとして取得することで、完全に自動化できます。ラベル付きデータセットは、モデルが学習するための「ラベル」を生成します。から。したがって、自己教師あり学習の最初の重要なポイントは、ラベルのないデータセットの教師あり情報をどのようにマイニングするかということです。自己教師あり学習では、教師なしモデルを学習するために教師あり手法を使用するため、教師あり情報をマイニングした後、その情報をどのように使用するかを検討する必要があります。これが自己教師あり学習の 2 番目の重要なポイント、つまり合理的なプロキシを設計する方法です。データ内の潜在的な特徴をマイニングするタスク。
対照学習は、最初の問題の解決策を提供します。画像にさまざまな強化 (トリミング、ノイズの追加など) が施された後、これらの強化された画像は正のサンプル ペアと見なされます。
Moco は、比較学習における最も古典的なモデルの 1 つです。これは、正のサンプルと負のサンプルに基づく比較学習手法です。正のサンプルと負のサンプルに基づく比較学習アルゴリズムでは、通常、負のサンプルはサンプル ライブラリ内の他の画像になりますが、これはサンプル ライブラリは同じではありませんが、これらのモデルは、マッピング空間内で陽性サンプルを十分近くに配置し、マッピング空間内で陰性サンプルを十分に遠くに配置するための異なる方法を試みており、これは個人識別の代理タスクです。サンプルがあるのでクラス。MOCO は、正サンプルと負サンプルに基づく比較学習を辞書引き問題としてまとめており、この問題の重要なポイントは、大規模で一貫性のある辞書を生成する方法です。
まず辞書クエリ問題について説明します 辞書内のデータはキー(key)と値(value)で構成されています 比較学習では絵をキーとして想像し、絵に対応する潜在的な特徴を値として比較します特定の画像がこの辞書でクエリされ、それが同じ画像からのものであるが、異なる拡張が施された後、一致が成功した場合、それらの値は可能な限り近くなる必要があります。問題は 2 つの部分に分かれています: 大きい部分と一貫性のある部分 以前の研究では、最初の部分は大きいです。この辞書は小さすぎる (ミニバッチ) か大きすぎる (データセット全体) ため、この 2 つの間のトレードオフは次のとおりです。 Moco のソリューションはサンプル キューであり、固定数のミニバッチがこのキューに保存され、新しいミニバッチがサンプル キューに追加されるたびに、最も古いミニバッチがデキューされます。その後、一貫性: 比較モデル学習は常に更新されるため、各サンプルが異なる時点でモデルを通過した後、取得される特徴は一貫性がないため、異なるエポックでは、同じサンプルのトレーニングに使用される特徴も異なります。MOCO の解は運動量です。エンコーダ、つまり運動量更新辞書 中央のサンプルの特性。その大部分は前のラウンドのトレーニングから来ています (MOCO の実験では、特徴の 99.9% が前のラウンドから来ていることが証明されており、効果はより優れています) )、加えて、サンプル キューは毎回最も古いサンプルを削除します。最も古いサンプルは最も勢いのある更新と最も高い不一致を持ちます。この方法により、サンプル キュー内の特徴の一貫性を確保できます。
2、コードを集中的に読む
2.1 コード構造
コードは 2 つのフォルダーといくつかのファイルに分割されています。フォルダー検出はターゲット検出の下流タスク用です。フォルダーmocoはモデルの主要部分です。ビルダーはメイン フォルダーの下にあります。main_moco は自己教師ありトレーニング プロセスです。 moco の、main_lincls は、画像分類タスク用の単純な線形分類器をトレーニングすることです。
読み取りプロセスは次のとおりです:
main_moco.py->moco フォルダー->main_cls.py (->検出フォルダー)
2.2 main_moco.py
2.2.1 パラメータの設定
model_names = sorted(name for name in models.__dict__
if name.islower() and not name.startswith("__")
and callable(models.__dict__[name]))
model_names は、torch 内のさまざまな視覚的なバックボーン ネットワーク名であり、さまざまなパラメータの意味は次のとおりです。
| パラメータ名 | パラメータの意味 |
|---|---|
| データ | データセットのパス |
| アーチ | バックボーン ネットワーク、model_name のいずれかを選択してください |
| 労働者 | データローダーのパラメータ |
| エポック | トレーニング |
| 開始エポック | 初期エポックは通常 0 です。特定のエポックの操作が中断された場合、このパラメーターを有効にしてトレーニングを続行できます。 |
| バッチサイズ | |
| 学習率 | モデルの学習率 |
| 勢い | モデルの勢い |
| 体重減少 | 体重減少 |
| 履歴書 | 最新のチェックポイントのパス |
| 粘液の薄暗い | 出力寸法 |
| モコク | サンプルキューサイズ (負のサンプルサイズ) |
| モコム | 辞書更新の勢い |
| モコちゃん | ソフトマックス温度 |
def main():
main は、最初に argparser によって渡されたさまざまなパラメーターを処理し、最後に main_worker 関数を最後に呼び出します。
main_worker(args.gpu, ngpus_per_node, args)
ここですべてデフォルトの場合、args.gpu=None、ngpus_per_node=GPU の数。
def main_worker(gpu, ngpus_per_node, args)
この関数は最初に複数のプロセスを処理し、次にパラメーターで選択されたバックボーン ネットワークと moco のいくつかの特別なパラメーターに従ってモデルを構築します。
print("=> creating model '{}'".format(args.arch))
model = moco.builder.MoCo(
models.__dict__[args.arch],
args.moco_dim, args.moco_k, args.moco_m, args.moco_t, args.mlp)
print(model)
モデルがクロスエントロピー損失関数を使用し、確率的勾配降下法を使用してモデルを最適化していることがわかります。
# 模型
model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu])
# 损失函数
criterion = nn.CrossEntropyLoss().cuda(args.gpu)
# 优化器
optimizer = torch.optim.SGD(model.parameters(), args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
次に、データ処理です。これは、データ強化の方法とデータの標準化の方法を定義します。
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
augmentation = [
transforms.RandomResizedCrop(224, scale=(0.2, 1.)),
transforms.RandomGrayscale(p=0.2),
transforms.ColorJitter(0.4, 0.4, 0.4, 0.4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize
]
次に、その後のトレーニングを容易にするためにデータセットとデータローダーが定義されます。train_dataset 内の各データは、異なる強化後の同じ画像のサンプル ペアです。
train_dataset = datasets.ImageFolder(
traindir,
moco.loader.TwoCropsTransform(transforms.Compose(augmentation)))
train_loader = torch.utils.data.DataLoader(
train_dataset, batch_size=args.batch_size, shuffle=(train_sampler is None),
num_workers=args.workers, pin_memory=True, sampler=train_sampler, drop_last=True)
最後に、上で定義した基本コンポーネントに従ってトレーニング プロセスを開始します。
train(train_loader, model, criterion, optimizer, epoch, args)
def train(train_loader、モデル、基準、オプティマイザ、エポック、引数)
ここでは、複数の処理・時間制御など、モデルの学習に関係のない処理をクリアしていきます。
for i, (images, _) in enumerate(train_loader):
# compute output
output, target = model(im_q=images[0], im_k=images[1])
loss = criterion(output, target)
losses.update(loss.item(), images[0].size(0))
# compute gradient and do SGD step
optimizer.zero_grad()
loss.backward()
optimizer.step()
コードに示されているように、モデルには 2 つのデータ入力があり、1 つは辞書クエリのキューに対応し、もう 1 つはキューに一致する辞書内のキーに対応します。
2.3 mocoフォルダ
mocoフォルダ内のmocoのモデル構造とデータ出力
2.3.1 ローダー.py
loader.py は非常にシンプルで、セクション 2.2.1 の main_worker で定義されたデータ拡張メソッドに従って、同じ画像の 2 つのブランチに異なる拡張を実装します。
class TwoCropsTransform:
"""Take two random crops of one image as the query and key."""
def __init__(self, base_transform):
self.base_transform = base_transform
def __call__(self, x):
q = self.base_transform(x)
k = self.base_transform(x)
return [q, k]
このうち、base_transform は main_worker で定義された拡張です
2.3.2 builder.py
モデルの初期化
def __init__(self, base_encoder, dim=128, K=65536, m=0.999, T=0.07, mlp=False):
"""
dim: feature dimension (default: 128)
K: queue size; number of negative keys (default: 65536)
m: moco momentum of updating key encoder (default: 0.999)
T: softmax temperature (default: 0.07)
"""
super(MoCo, self).__init__()
# 基本参数
self.K = K
self.m = m
self.T = T
# create the encoders
# num_classes is the output fc dimension
self.encoder_q = base_encoder(num_classes=dim)
self.encoder_k = base_encoder(num_classes=dim)
if mlp: # hack: brute-force replacement
dim_mlp = self.encoder_q.fc.weight.shape[1]
self.encoder_q.fc = nn.Sequential(nn.Linear(dim_mlp, dim_mlp), nn.ReLU(), self.encoder_q.fc)
self.encoder_k.fc = nn.Sequential(nn.Linear(dim_mlp, dim_mlp), nn.ReLU(), self.encoder_k.fc)
for param_q, param_k in zip(self.encoder_q.parameters(), self.encoder_k.parameters()):
param_k.data.copy_(param_q.data) # initialize
param_k.requires_grad = False # not update by gradient
# create the queue
self.register_buffer("queue", torch.randn(dim, K))
self.queue = nn.functional.normalize(self.queue, dim=0)
self.register_buffer("queue_ptr", torch.zeros(1, dtype=torch.long))
ここに保存されているのは moco のモデルです まず、キューのモデルとキーのモデルは構造的に完全に一致していることがわかりますが、キーのモデルは勾配リターンを持たず、直接論文の設計とは異なるキュー モデルからコピーされたものであり、一貫性があること。
サンプルキュー
初期化の最後に、モデルはサンプル キューも定義します。ここで、キューはキュー ヘッドとリストによって維持される循環キューであり、キューのエンキュー/終了操作は次の関数で示されます。
def _dequeue_and_enqueue(self, keys):
# gather keys before updating queue
keys = concat_all_gather(keys)
batch_size = keys.shape[0]
ptr = int(self.queue_ptr)
assert self.K % batch_size == 0 # for simplicity
# replace the keys at ptr (dequeue and enqueue)
self.queue[:, ptr:ptr + batch_size] = keys.T
ptr = (ptr + batch_size) % self.K # move pointer
self.queue_ptr[0] = ptr
毎回、キューはキュー ヘッドが指すバッチを新しいバッチに置き換え、次にキュー ヘッドが新しく追加されたバッチの末尾を指すようにします。これは、最も古いバッチをキューに入れて、新しいバッチをキューに追加するのと同じです。しっぽ。
運動量エンコーダ
def _momentum_update_key_encoder(self):
"""
Momentum update of the key encoder
"""
for param_q, param_k in zip(self.encoder_q.parameters(), self.encoder_k.parameters()):
param_k.data = param_k.data * self.m + param_q.data * (1. - self.m)
param_q は勾配を介して戻される更新されたパラメーター、param_k は上記のループ操作を実行する前の前の特徴、self.m は運動量、最良は 0.999、つまり、サンプル ディクショナリ内の特徴の 99.9% は次のとおりです。以前の機能のうち、現在の更新によるものはわずか 0.1% であり、辞書の高い一貫性が保証されます。
モデルフォワードプロセス
転送プロセスは比較的従来通りで、まず、キューのイメージがキュー エンコーダに渡されて、次の特徴が取得されます。
# compute query features
q = self.encoder_q(im_q) # queries: NxC
q = nn.functional.normalize(q, dim=1)
次に、勢いによってキー エンコーダーが更新され、キーの機能が取得されます。
self._momentum_update_key_encoder() # update the key encoder
k = self.encoder_k(im_k) # keys: NxC
k = nn.functional.normalize(k, dim=1)
次に、キューとキーに従ってモデルの損失を取得します。ここで、q と k は互いの正の例であるため、エラーは l_pos と呼ばれます。この時点では新しいサンプルはまだキューに入っていないため、q とすべてサンプル キュー内のサンプルは互いに負の例であるため、q とサンプル キューの間の誤差は l_neg と呼ばれます。
l_pos = torch.einsum('nc,nc->n', [q, k]).unsqueeze(-1)
l_neg = torch.einsum('nc,ck->nk', [q, self.queue.clone().detach()])
次に、torch.cat を使用して正の誤差と負の誤差を結合します。ここでは l_pos が最初に来るので、正のサンプル誤差は常に各行の 0 番目の要素になることに注意してください。
# logits: Nx(1+K)
logits = torch.cat([l_pos, l_neg], dim=1)
# apply temperature
logits /= self.T
正のサンプルは各行の 0 番目の要素にあるため、クロス エントロピーを計算する場合、入力ラベルは正のサンプルの位置を表すため、ラベルはすべて 0 になります。
# labels: positive key indicators
labels = torch.zeros(logits.shape[0], dtype=torch.long).cuda()
最後に、サンプル キューを更新し、古いサンプルをキューから取り出し、新しいサンプルをキューに追加します。
# dequeue and enqueue
self._dequeue_and_enqueue(k)
エラーとラベルを返した後、main_moco.pyのmain_worker関数でクロスエントロピーロスが計算され、キューエンコーダのパラメータが更新されます。
return logits, labels
上記は完全な事前トレーニング プロセスを構成します。
2.4 main_cls.py
この部分は主に、事前トレーニング済みモデルを使用し、ダウンストリーム タスクで微調整してテストすることです。メイン プロセスは main_moco と似ていますが、違いは次のとおりです。
モデル構造は異なります。
事前トレーニングされたキュー エンコーダーは main_cls で直接抽出されます。
'''
首先，构建骨干网络实例
'''
model = models.__dict__[args.arch]()
'''
然后，加载预训练模型，并保留queue编码器部分
'''
checkpoint = torch.load(args.pretrained, map_location="cpu")
# rename moco pre-trained keys
state_dict = checkpoint['state_dict']
for k in list(state_dict.keys()):
# retain only encoder_q up to before the embedding layer
if k.startswith('module.encoder_q') and not k.startswith('module.encoder_q.fc'):
state_dict[k[len("module.encoder_q."):]] = state_dict[k]
# delete renamed or unused k
del state_dict[k]
'''
最后，将保留的queue编码器加载到构建的实例中
'''
msg = model.load_state_dict(state_dict, strict=False)
使用部分は異なります。
この部分は主にモデルの微調整のためのものであるため、トレーニングのようにモデルのすべての層に対する勾配復帰は実行されませんが、モデルの線形層のみが更新されるため、次のような他の層は更新されません。 CNN 層をフリーズする必要があります。
# freeze all layers but the last fc
for name, param in model.named_parameters():
if name not in ['fc.weight', 'fc.bias']:
param.requires_grad = False
# init the fc layer
model.fc.weight.data.normal_(mean=0.0, std=0.01)
model.fc.bias.data.zero_()