畳み込みオートエンコーダ
畳み込みニューラルネットワークは、2次元の空間構造データインスタンス（画像）を取り込み、それを処理して何らかの1次元ベクトル表現を生成します。こうした背景から、画像マトリックスからベクトル表現へのマッピングが学習できるのであれば、そのベクトル表現から画像へのマッピングも学習できるのではないかという疑問が生まれます。このデモ記事では、まさにそのことを探求していきます。
前提条件
このチュートリアルを進めるためには、Pythonコードとニューラルネットワークに関する基本的な理解が必要です。新しいアーキテクチャの開発経験がある中級者から上級者向けのコーダーに、この資料をお勧めします。
この記事のコードは、通常の家庭用PCまたはDigitalOceanのDropletで実行できます。
特徴抽出器としての畳み込みニューラルネットワーク
以前の記事では、畳み込みネットワークの畳み込み層が画像から特徴を抽出する目的であることに触れました。それらの特徴は、その後、実際の分類を行う線形層に渡されます（1 x 1 の畳み込み層をダウンサンプリングに利用するアーキテクチャに対しては例外があります）。
上記に示されたアーキテクチャのVGG-16を考慮すると、入力層からプールされた7 x 7 x 512特徴マップがフラット化されて25,088要素のベクターが作成されるポイントまでの部分が特徴抽出器として機能します。つまり、合計50,176ピクセルの224 x 224画像が処理されて25,088要素の特徴ベクターが作成され、この特徴ベクターはその後、分類のために線形層に渡されます。
これらの特徴は畳み込みネットワークによって抽出されるため、別の畳み込みネットワークがこれらの特徴を理解し、それらの特徴が属する元の画像を再構築することができると考えるのは論理的です。これは基本的に特徴抽出プロセスを逆転させることになります。これが基本的にオートエンコーダーの役割です。
オートエンコーダーの構造
前のセクションで述べたように、オートエンコーダは特徴ベクトルからデータインスタンスを再構成することができる深層学習アーキテクチャです。彼らはあらゆる種類のデータで機能しますが、この文書は主に画像データへの応用に関して焦点を当てています。オートエンコーダは、主に3つのコンポーネントで構成されています。すなわち、エンコーダ、ボトルネック、およびデコーダです。
エンコーダー
オートエンコーダの最初のセクションであるエンコーダは、特徴抽出器として特に機能する畳み込みネットワークです。その主な機能は、画像から最も顕著な特徴を抽出し、それをベクトルとして返すことです。
ボトルネック
エンコーダの直後に位置するボトルネックは、コード層とも呼ばれ、抽出された特徴をより小さなベクトル表現に圧縮するのを助ける追加の層として機能します。これは、デコーダが特徴を理解するのを難しくし、より複雑なマッピングを学ぶよう強制するために行われます。
デコーダー
オートエンコーダの最後の部分であるデコーダは、エンコーダから出てくる特徴を把握しようとする畳み込みネットワークであり、ボトルネックで圧縮された特徴を使用して、元の画像をそのまま再構築することを目指しています。
オートエンコーダの訓練
このセクションでは、PyTorchでオートエンコーダをゼロから実装し、特定のデータセットでトレーニングします。
必要なパッケージをすぐにインポートすることから始めましょう。
article dependencies
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision
import torchvision.transforms as transforms
import torchvision.datasets as Datasets
from torch.utils.data import Dataset, DataLoader
import numpy as np
import matplotlib.pyplot as plt
import cv2
from tqdm.notebook import tqdm
from tqdm import tqdm as tqdm_regular
import seaborn as sns
from torchvision.utils import make_grid
import random
configuring device
if torch.cuda.is_available():
device = torch.device('cuda:0')
print('Running on the GPU')
else:
device = torch.device('cpu')
print('Running on the CPU')
データの準備
この記事の目的のために、畳み込みオートエンコーダのトレーニングにCIFAR-10データセットを使用します。以下のコードセルでのようにロードできます。
loading training data
training_set = Datasets.CIFAR10(root='./', download=True,
transform=transforms.ToTensor())
loading validation data
validation_set = Datasets.CIFAR10(root='./', download=True, train=False,
transform=transforms.ToTensor())
次に、データセットから画像のみを抽出する必要があります。オートエンコーダに画像を再構成させようとしているため、ターゲットはクラスラベルではなく、実際の画像そのものになります。また、各クラスから1枚の画像を抽出し、視覚化の目的でオブジェクト「test_images」に保存します。これについては後ほど詳しく説明します。
def extract_each_class(dataset):
"""
This function searches for and returns
one image per class
"""
images = []
ITERATE = True
i = 0
j = 0
while ITERATE:
for label in tqdm_regular(dataset.targets):
if label==j:
images.append(dataset.data[i])
print(f'class {j} found')
i+=1
j+=1
if j==10:
ITERATE = False
else:
i+=1
return images
extracting training images
training_images = [x for x in training_set.data]
extracting validation images
validation_images = [x for x in validation_set.data]
extracting test images for visualization purposes
test_images = extract_each_class(validation_set)
さて、画像をテンソルとして使用できるようにするために、PyTorch のデータセットクラスを定義する必要があります。これとクラスのインスタンス化は、以下のコードセルで行われます。
defining dataset class
class CustomCIFAR10(Dataset):
def __init__(self, data, transforms=None):
self.data = data
self.transforms = transforms
def __len__(self):
return len(self.data)
def __getitem__(self, idx):
image = self.data[idx]
if self.transforms!=None:
image = self.transforms(image)
return image
creating pytorch datasets
training_data = CustomCIFAR10(training_images, transforms=transforms.Compose([transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]))
validation_data = CustomCIFAR10(validation_images, transforms=transforms.Compose([transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]))
test_data = CustomCIFAR10(test_images, transforms=transforms.Compose([transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]))
オートエンコーダーアーキテクチャ
この記事の目的のためにカスタム畳み込みオートエンコーダーアーキテクチャが以下のように定義されます。このアーキテクチャはCIFAR-10データセットで動作するように設計されており、エンコーダーは32 x 32ピクセルの画像を3チャンネルで受け取り、64の8 x 8特徴マップが生成されるまで処理します。
これらの特徴マップは、その後、4096要素のベクターにフラット化され、ボトルネックでわずか200要素に圧縮されます。デコーダーはこの200要素のベクター表現を受け取り、転置畳み込みを通じて処理し、3 x 32 x 32の画像が出力として返されます。
上で定義されたアーキテクチャは、以下のコードセルに実装されています。この場合のパラメータ「latent_dim」は、ボトルネックのサイズを指し、200として指定しています。
defining encoder
class Encoder(nn.Module):
def __init__(self, in_channels=3, out_channels=16, latent_dim=200, act_fn=nn.ReLU()):
super().__init__()
self.net = nn.Sequential(
nn.Conv2d(in_channels, out_channels, 3, padding=1), # (32, 32)
act_fn,
nn.Conv2d(out_channels, out_channels, 3, padding=1),
act_fn,
nn.Conv2d(out_channels, 2*out_channels, 3, padding=1, stride=2), # (16, 16)
act_fn,
nn.Conv2d(2*out_channels, 2*out_channels, 3, padding=1),
act_fn,
nn.Conv2d(2*out_channels, 4*out_channels, 3, padding=1, stride=2), # (8, 8)
act_fn,
nn.Conv2d(4*out_channels, 4*out_channels, 3, padding=1),
act_fn,
nn.Flatten(),
nn.Linear(4*out_channels*8*8, latent_dim),
act_fn
)
def forward(self, x):
x = x.view(-1, 3, 32, 32)
output = self.net(x)
return output
defining decoder
class Decoder(nn.Module):
def __init__(self, in_channels=3, out_channels=16, latent_dim=200, act_fn=nn.ReLU()):
super().__init__()
self.out_channels = out_channels
self.linear = nn.Sequential(
nn.Linear(latent_dim, 4*out_channels*8*8),
act_fn
)
self.conv = nn.Sequential(
nn.ConvTranspose2d(4*out_channels, 4*out_channels, 3, padding=1), # (8, 8)
act_fn,
nn.ConvTranspose2d(4*out_channels, 2*out_channels, 3, padding=1,
stride=2, output_padding=1), # (16, 16)
act_fn,
nn.ConvTranspose2d(2*out_channels, 2*out_channels, 3, padding=1),
act_fn,
nn.ConvTranspose2d(2*out_channels, out_channels, 3, padding=1,
stride=2, output_padding=1), # (32, 32)
act_fn,
nn.ConvTranspose2d(out_channels, out_channels, 3, padding=1),
act_fn,
nn.ConvTranspose2d(out_channels, in_channels, 3, padding=1)
)
def forward(self, x):
output = self.linear(x)
output = output.view(-1, 4*self.out_channels, 8, 8)
output = self.conv(output)
return output
defining autoencoder
class Autoencoder(nn.Module):
def __init__(self, encoder, decoder):
super().__init__()
self.encoder = encoder
self.encoder.to(device)
self.decoder = decoder
self.decoder.to(device)
def forward(self, x):
encoded = self.encoder(x)
decoded = self.decoder(encoded)
return decoded
いつものように、トレーニングとバリデーションをよりシームレスにするためのクラスを定義する必要があります。この場合、生成モデルをトレーニングしているため、損失はあまり多くの情報を持たないかもしれません。一般的に、損失を減少させたいと考えており、損失値を使用して各エポックごとにオートエンコーダが画像を再構築する能力を確認することもできます。このため、以下に示すように視覚化ブロックを含めました。
class ConvolutionalAutoencoder():
def __init__(self, autoencoder):
self.network = autoencoder
self.optimizer = torch.optim.Adam(self.network.parameters(), lr=1e-3)
def train(self, loss_function, epochs, batch_size,
training_set, validation_set, test_set):
# creating log
log_dict = {
'training_loss_per_batch': [],
'validation_loss_per_batch': [],
'visualizations': []
}
# defining weight initialization function
def init_weights(module):
if isinstance(module, nn.Conv2d):
torch.nn.init.xavier_uniform_(module.weight)
module.bias.data.fill_(0.01)
elif isinstance(module, nn.Linear):
torch.nn.init.xavier_uniform_(module.weight)
module.bias.data.fill_(0.01)
# initializing network weights
self.network.apply(init_weights)
# creating dataloaders
train_loader = DataLoader(training_set, batch_size)
val_loader = DataLoader(validation_set, batch_size)
test_loader = DataLoader(test_set, 10)
# setting convnet to training mode
self.network.train()
self.network.to(device)
for epoch in range(epochs):
print(f'Epoch {epoch+1}/{epochs}')
train_losses = []
#------------
# TRAINING
#------------
print('training...')
for images in tqdm(train_loader):
# zeroing gradients
self.optimizer.zero_grad()
# sending images to device
images = images.to(device)
# reconstructing images
output = self.network(images)
# computing loss
loss = loss_function(output, images.view(-1, 3, 32, 32))
# calculating gradients
loss.backward()
# optimizing weights
self.optimizer.step()
#--------------
# LOGGING
#--------------
log_dict['training_loss_per_batch'].append(loss.item())
#--------------
# VALIDATION
#--------------
print('validating...')
for val_images in tqdm(val_loader):
with torch.no_grad():
# sending validation images to device
val_images = val_images.to(device)
# reconstructing images
output = self.network(val_images)
# computing validation loss
val_loss = loss_function(output, val_images.view(-1, 3, 32, 32))
#--------------
# LOGGING
#--------------
log_dict['validation_loss_per_batch'].append(val_loss.item())
#--------------
# VISUALISATION
#--------------
print(f'training_loss: {round(loss.item(), 4)} validation_loss: {round(val_loss.item(), 4)}')
for test_images in test_loader:
# sending test images to device
test_images = test_images.to(device)
with torch.no_grad():
# reconstructing test images
reconstructed_imgs = self.network(test_images)
# sending reconstructed and images to cpu to allow for visualization
reconstructed_imgs = reconstructed_imgs.cpu()
test_images = test_images.cpu()
# visualisation
imgs = torch.stack([test_images.view(-1, 3, 32, 32), reconstructed_imgs],
dim=1).flatten(0,1)
grid = make_grid(imgs, nrow=10, normalize=True, padding=1)
grid = grid.permute(1, 2, 0)
plt.figure(dpi=170)
plt.title('Original/Reconstructed')
plt.imshow(grid)
log_dict['visualizations'].append(grid)
plt.axis('off')
plt.show()
return log_dict
def autoencode(self, x):
return self.network(x)
def encode(self, x):
encoder = self.network.encoder
return encoder(x)
def decode(self, x):
decoder = self.network.decoder
return decoder(x)
すべてが設定されたので、次に以下で定義した畳み込みオートエンコーダークラスのメンバーとして、続くコードセルで指定したパラメータを使用して、オートエンコーダーをインスタンス化できます。
training model
model = ConvolutionalAutoencoder(Autoencoder(Encoder(), Decoder()))
log_dict = model.train(nn.MSELoss(), epochs=10, batch_size=64,
training_set=training_data, validation_set=validation_data,
test_set=test_data)
最初のエポックの終わりから、私たちのデコーダーがエンコーダーに供給された画像を再構築する方法を発達させ始めたことは明らかです。エンコーダーは圧縮された200要素の特徴ベクトル表現しか利用できなかったにもかかわらず、10エポック目まで再構築された画像は詳細が増え続けています。
トレーニングと検証の損失を見ていると、オートエンコーダはまだトレーニングのエポックを少し追加することで利益を得る可能性がある。というのも、その損失はまだ減少傾向にあるからだ。これはトレーニングの損失よりも検証の損失に当てはまるようで、トレーニングの損失は停滞しているように見える。
ボトルネックと詳細
以前のセクションの一つで、ボトルネックコーディングレイヤーが特徴ベクトルをさらに圧縮する目的を果たすことについて言及しました。その結果、デコーダーがより複雑で一般化可能なマッピングを学習することを促します。一方で、コーディングレイヤーの圧縮の度合いがデコーダーが画像を再構築する能力に影響を与えるため、微妙なバランスを見つける必要があります。
デコーダに渡されるベクトル表現が小さいほど、デコーダがアクセスできる画像の特徴が少なくなり、その再構築の詳細も少なくなります。同様に、デコーダに渡されるベクトル表現が大きいほど、アクセスできる画像の特徴が増え、その再構築の詳細も増えます。この考えに従って、同じオートエンコーダのアーキテクチャを訓練しますが、今回はサイズ1000のボトルネックを使用します。
training model
model = ConvolutionalAutoencoder(Autoencoder(Encoder(latent_dim=1000), Decoder(latent_dim=1000)))
log_dict = model.train(nn.MSELoss(), epochs=10, batch_size=64,
training_set=training_data, validation_set=validation_data,
test_set=test_data)
各エポックごとに生成された視覚化から、デコーダーが画像の詳細と視覚的精度の再構築において優れた仕事をしていることがすぐに明らかです。これは、新しいデコーダーがより多くの特徴にアクセスできるという事実に起因しています。元の4096要素の特徴ベクトルは、200ではなく1000要素にダウンサンプリングされています。
エポック1（上）対エポック10（下）。
再度、オートエンコーダはもう少しエポックのトレーニングを受けることで利益を得られるかもしれません。それのトレーニングおよび検証損失はまだ下降傾向にあり、私たちがボトルネックを200要素のみにしてオートエンコーダをトレーニングしたときに観察した傾斜よりも急な傾斜を持っています。
サイズ200と1000のボトルネックを10エポック目で比較すると、1000要素のボトルネックで生成された画像は、200要素のボトルネックで生成された画像よりも明確で詳細であることが明らかになります。
ボトルネック200（上）対ボトルネック1000（下）、どちらも10エポック目におけるもの。
絶対的な洗練への訓練
畳み込みオートエンコーダーはどの時点で最適にトレーニングされるのでしょうか？私たちがトレーニングした2つのオートエンコーダーから、損失プロットが平坦になり始めたにもかかわらず、10エポック目では再構成がまだぼやけていることが観察できます。ボトルネックのサイズを増やすことはこの問題をある程度改善するだけですが、完全には解決しません。
これは、今回使用された損失関数、平均二乗誤差に部分的に起因しています。これは生成モデルの損失を測定する際にあまり効果的ではありません。ほとんどの場合、このようなぼやけた再構成は畳み込みオートエンコーダの作業にとっては厄介なものです。もし再構成したり画像を生成することが目標であれば、敵対的生成ネットワーク（GAN）や拡散モデルの方が安全な選択かもしれません。しかし、畳み込みオートエンコーダが無用だと言うわけではなく、異常検出や画像ノイズ除去などに利用できます。
最後の remarks
この記事では、画像データの文脈におけるオートエンコーダーについて説明しました。次に、畳み込みオートエンコーダーが具体的に何を行うのか、そしてその動作原理の直感を養うためにどのように行うのかを見ていきました。その後、さまざまなセクションに触れ、自分たちのカスタムオートエンコーダーを定義し、それを訓練し、モデル訓練の結果について議論しました。