在PyTorch中使用深度自編碼器實現(xiàn)圖像重建
點擊上方“小白學(xué)視覺”，選擇加"星標(biāo)"或“置頂”
重磅干貨，第一時間送達(dá)
? 磐創(chuàng)AI分享??
? 磐創(chuàng)AI分享??
作者 | DR. VAIBHAV KUMAR?
編譯 | VK?
來源 | Analytics In Diamag
人工神經(jīng)網(wǎng)絡(luò)有許多流行的變體，可用于有監(jiān)督和無監(jiān)督學(xué)習(xí)問題。自編碼器也是神經(jīng)網(wǎng)絡(luò)的一個變種，主要用于無監(jiān)督學(xué)習(xí)問題。
當(dāng)它們在體系結(jié)構(gòu)中有多個隱藏層時，它們被稱為深度自編碼器。這些模型可以應(yīng)用于包括圖像重建在內(nèi)的各種應(yīng)用。
在圖像重建中，他們學(xué)習(xí)輸入圖像模式的表示，并重建與原始輸入圖像模式匹配的新圖像。圖像重建有許多重要的應(yīng)用，特別是在醫(yī)學(xué)領(lǐng)域，需要從現(xiàn)有的不完整或有噪聲的圖像中提取解碼后的無噪聲圖像。
在本文中，我們將演示在PyTorch中實現(xiàn)用于重建圖像的深度自編碼器。該深度學(xué)習(xí)模型將以MNIST手寫數(shù)字為訓(xùn)練對象，在學(xué)習(xí)輸入圖像的表示后重建數(shù)字圖像。
自編碼器
自編碼器是人工神經(jīng)網(wǎng)絡(luò)的變體，通常用于以無監(jiān)督的方式學(xué)習(xí)有效的數(shù)據(jù)編碼。
他們通常在一個表示學(xué)習(xí)方案中學(xué)習(xí)，在那里他們學(xué)習(xí)一組數(shù)據(jù)的編碼。網(wǎng)絡(luò)通過學(xué)習(xí)輸入數(shù)據(jù)的表示，以非常相似的方式重建輸入數(shù)據(jù)。自編碼器的基本結(jié)構(gòu)如下所示。
該體系結(jié)構(gòu)通常包括輸入層、輸出層和連接輸入和輸出層的一個或多個隱藏層。輸出層與輸入層具有相同數(shù)量的節(jié)點，因為它要重新構(gòu)造輸入。
在它的一般形式中，只有一個隱藏層，但在深度自動編碼器的情況下，有多個隱藏層。這種深度的增加減少了表示某些函數(shù)的計算成本，也減少了學(xué)習(xí)某些函數(shù)所需的訓(xùn)練數(shù)據(jù)量。其應(yīng)用領(lǐng)域包括異常檢測、圖像處理、信息檢索、藥物發(fā)現(xiàn)等。
在PyTorch中實現(xiàn)深度自編碼器
首先，我們將導(dǎo)入所有必需的庫。
import?os
import?torch?
import?torchvision
import?torch.nn?as?nn
import?torchvision.transforms?as?transforms
import?torch.optim?as?optim
import?matplotlib.pyplot?as?plt
import?torch.nn.functional?as?F
from?torchvision?import?datasets
from?torch.utils.data?import?DataLoader
from?torchvision.utils?import?save_image
from?PIL?import?Image
現(xiàn)在，我們將定義超參數(shù)的值。
Epochs?=?100
Lr_Rate?=?1e-3
Batch_Size?=?128
以下函數(shù)將用于PyTorch模型所需的圖像轉(zhuǎn)換。
transform?=?transforms.Compose([
????transforms.ToTensor(),
????transforms.Normalize((0.5,),?(0.5,))
])
使用下面的代碼片段，我們將下載MNIST手寫數(shù)字?jǐn)?shù)據(jù)集，并為進(jìn)一步處理做好準(zhǔn)備。
train_set?=?datasets.MNIST(root='./data',?train=True,?download=True,?transform=transform)
test_set?=?datasets.MNIST(root='./data',?train=False,?download=True,?transform=transform)
train_loader?=?DataLoader(train_set,?Batch_Size=Batch_Size,?shuffle=True)
test_loader?=?DataLoader(test_set,?Batch_Size=Batch_Size,?shuffle=True)
讓我們看看關(guān)于訓(xùn)練數(shù)據(jù)及其類的一些信息。
print(train_set)
print(train_set.classes)
在下一步中，我們將定義用于定義模型的Autoencoder類。
class?Autoencoder(nn.Module):
????def?__init__(self):
????????super(Autoencoder,?self).__init__()
????????#編碼器
????????self.enc1?=?nn.Linear(in_features=784,?out_features=256)?#?Input?image?(28*28?=?784)
????????self.enc2?=?nn.Linear(in_features=256,?out_features=128)
????????self.enc3?=?nn.Linear(in_features=128,?out_features=64)
????????self.enc4?=?nn.Linear(in_features=64,?out_features=32)
????????self.enc5?=?nn.Linear(in_features=32,?out_features=16)
????????#解碼器?
????????self.dec1?=?nn.Linear(in_features=16,?out_features=32)
????????self.dec2?=?nn.Linear(in_features=32,?out_features=64)
????????self.dec3?=?nn.Linear(in_features=64,?out_features=128)
????????self.dec4?=?nn.Linear(in_features=128,?out_features=256)
????????self.dec5?=?nn.Linear(in_features=256,?out_features=784)?#?Output?image?(28*28?=?784)
????def?forward(self,?x):
????????x?=?F.relu(self.enc1(x))
????????x?=?F.relu(self.enc2(x))
????????x?=?F.relu(self.enc3(x))
????????x?=?F.relu(self.enc4(x))
????????x?=?F.relu(self.enc5(x))
????????x?=?F.relu(self.dec1(x))
????????x?=?F.relu(self.dec2(x))
????????x?=?F.relu(self.dec3(x))
????????x?=?F.relu(self.dec4(x))
????????x?=?F.relu(self.dec5(x))
????????return?x
現(xiàn)在，我們將創(chuàng)建Autoencoder模型作為上面定義的Autoencoder類的一個對象。
model?=?Autoencoder()
print(model)
現(xiàn)在，我們將定義損失函數(shù)和優(yōu)化方法。
criterion?=?nn.MSELoss()
optimizer?=?optim.Adam(net.parameters(),?lr=Lr_Rate)
以下函數(shù)將啟用CUDA環(huán)境。
def?get_device():
????if?torch.cuda.is_available():
????????device?=?'cuda:0'
????else:
????????device?=?'cpu'
????return?device
下面的函數(shù)將創(chuàng)建一個目錄來保存結(jié)果。
def?make_dir():
????image_dir?=?'MNIST_Out_Images'
????if?not?os.path.exists(image_dir):
????????os.makedirs(image_dir)
使用下面的函數(shù)，我們將保存模型生成的重建圖像。
def?save_decod_img(img,?epoch):
????img?=?img.view(img.size(0),?1,?28,?28)
????save_image(img,?'./MNIST_Out_Images/Autoencoder_image{}.png'.format(epoch))
將調(diào)用下面的函數(shù)來訓(xùn)練模型。
def?training(model,?train_loader,?Epochs):
????train_loss?=?[]
????for?epoch?in?range(Epochs):
????????running_loss?=?0.0
????????for?data?in?train_loader:
????????????img,?_?=?data
????????????img?=?img.to(device)
????????????img?=?img.view(img.size(0),?-1)
????????????optimizer.zero_grad()
????????????outputs?=?model(img)
????????????loss?=?criterion(outputs,?img)
????????????loss.backward()
????????????optimizer.step()
????????????running_loss?+=?loss.item()
????????loss?=?running_loss?/?len(train_loader)
????????train_loss.append(loss)
????????print('Epoch?{}?of?{},?Train?Loss:?{:.3f}'.format(
????????????epoch+1,?Epochs,?loss))
????????if?epoch?%?5?==?0:
????????????save_decod_img(outputs.cpu().data,?epoch)
????return?train_loss
以下函數(shù)將對訓(xùn)練后的模型進(jìn)行圖像重建測試。
def?test_image_reconstruct(model,?test_loader):
?????for?batch?in?test_loader:
????????img,?_?=?batch
????????img?=?img.to(device)
????????img?=?img.view(img.size(0),?-1)
????????outputs?=?model(img)
????????outputs?=?outputs.view(outputs.size(0),?1,?28,?28).cpu().data
????????save_image(outputs,?'MNIST_reconstruction.png')
????????break
在訓(xùn)練之前，模型將被推送到CUDA環(huán)境中，并使用上面定義的函數(shù)創(chuàng)建目錄來保存結(jié)果圖像。
device?=?get_device()
model.to(device)
make_dir()
現(xiàn)在，將對模型進(jìn)行訓(xùn)練。
train_loss?=?training(model,?train_loader,?Epochs)
訓(xùn)練成功后，我們將在訓(xùn)練中可視化損失。
plt.figure()
plt.plot(train_loss)
plt.title('Train?Loss')
plt.xlabel('Epochs')
plt.ylabel('Loss')
plt.savefig('deep_ae_mnist_loss.png')
我們將可視化訓(xùn)練期間保存的一些圖像。
Image.open('/content/MNIST_Out_Images/Autoencoder_image0.png')
Image.open('/content/MNIST_Out_Images/Autoencoder_image50.png')
Image.open('/content/MNIST_Out_Images/Autoencoder_image95.png')
在最后一步，我們將測試我們的自編碼器模型來重建圖像。
test_image_reconstruct(model,?testloader)
Image.open('/content/MNIST_reconstruction.png')
所以，我們可以看到，自訓(xùn)練過程開始時，自編碼器模型就開始重建圖像。第一個epoch以后，重建的質(zhì)量不是很好，直到50 epoch后才得到改進(jìn)。
經(jīng)過完整的訓(xùn)練，我們可以看到，在95 epoch以后生成的圖像和測試中，它可以構(gòu)造出與原始輸入圖像非常匹配的圖像。
我們根據(jù)loss值，可以知道epoch可以設(shè)置100或200。
經(jīng)過長時間的訓(xùn)練，有望獲得更清晰的重建圖像。然而，通過這個演示，我們可以理解如何在PyTorch中實現(xiàn)用于圖像重建的深度自編碼器。
參考文獻(xiàn)：
Sovit Ranjan Rath, “Implementing Deep Autoencoder in PyTorch” Abien Fred Agarap, “Implementing an Autoencoder in PyTorch” Reyhane Askari, “Auto Encoders”
原文鏈接：https://analyticsindiamag.com/hands-on-guide-to-implement-deep-autoencoder-in-pytorch-for-image-reconstruction/