【PyTorch】實作AutoEncoder

實作AutoEncoder

訓練模型

import torch
import torch.nn as nn
import torch.utils.data as data
import torchvision


# ------- Settings -------
epochs = 10 # 訓練迭代次數
batch_size = 128
lr = 0.008 # 學習率


# ------- Download Training DataSet -------
# 使用MNIST DataSet，download設為true，若沒有會自動下載
train_set = torchvision.datasets.MNIST(
    root='mnist', # 數據存放路徑
    train=True, # True: training data，False: testing data
    download=True, # True: 下載，False: 不下載
    transform=torchvision.transforms.ToTensor(), # 設定load dataset時需要進行哪種變換操作，此處為轉成tensor格式
)

# ------- DataLoader -------
train_loader = data.DataLoader(train_set, batch_size=batch_size, shuffle=True)


# ------- AutoEncoder Model Structure -------
class AutoEncoder(nn.Module):
    def __init__(self):
        super(AutoEncoder, self).__init__()

        # Encoder
        self.encoder = nn.Sequential(
            nn.Linear(784, 128),
            nn.Tanh(),
            nn.Linear(128, 64),
            nn.Tanh(),
            nn.Linear(64, 16),
            nn.Tanh(),
            nn.Linear(16, 2),
        )

        # Decoder
        self.decoder = nn.Sequential(
            nn.Linear(2, 16),
            nn.Tanh(),
            nn.Linear(16, 64),
            nn.Tanh(),
            nn.Linear(64, 128),
            nn.Tanh(),
            nn.Linear(128, 784),
            nn.Sigmoid()
        )

    def forward(self, inputs):
        z = self.encoder(inputs)
        outputs = self.decoder(z)

        return z, outputs


# ------- Optimizer and Loss function -------
model = AutoEncoder() # 初始化模型
print(model) # print出model架構
optimizer = torch.optim.Adam(model.parameters(), lr=lr) # 選擇優化器
loss_function = nn.MSELoss() # 選擇損失函數


# ------- Train -------
for epoch in range(epochs):
    for data, labels in train_loader:
        # data size [128, 1, 28, 28] => inputs size [128,784]
        inputs = data.view(-1, 784)

        # Forward
        z, outputs= model(inputs) # 計算預測值

        # Backward
        optimizer.zero_grad() # 將模型參數梯度初始化為0
        loss = loss_function(outputs, inputs) # 計算當前損失
        loss.backward() # 計算梯度
        optimizer.step() # 更新所有參數

    # Show progress
    print('[{}/{}] Loss:'.format(epoch+1, epochs), round(loss.item(),5))


# ------- Save -------
torch.save(model, 'autoencoder.pth')

• 模型架構
  
• 訓練過程
  

測試模型

import torch
import torch.nn as nn
import torch.utils.data as data
import torchvision
import numpy as np
import matplotlib.pyplot as plt


# ------- Settings -------
plt.rcParams['figure.figsize'] = (10.0, 8.0)
plt.rcParams['image.interpolation'] = 'nearest'
plt.rcParams['image.cmap'] = 'gray'


# ------- Show images -------
def show_images(images):
    sqrtn = int(np.ceil(np.sqrt(images.shape[0])))

    for index, image in enumerate(images):
        plt.subplot(sqrtn, sqrtn, index+1)
        plt.imshow(image.reshape(28, 28))
        plt.axis('off')


# ------- AutoEncoder Model Structure -------
class AutoEncoder(nn.Module):
    def __init__(self):
        super(AutoEncoder, self).__init__()

        # Encoder
        self.encoder = nn.Sequential(
            nn.Linear(784, 128),
            nn.Tanh(),
            nn.Linear(128, 64),
            nn.Tanh(),
            nn.Linear(64, 16),
            nn.Tanh(),
            nn.Linear(16, 2),
        )

        # Decoder
        self.decoder = nn.Sequential(
            nn.Linear(2, 16),
            nn.Tanh(),
            nn.Linear(16, 64),
            nn.Tanh(),
            nn.Linear(64, 128),
            nn.Tanh(),
            nn.Linear(128, 784),
            nn.Sigmoid()
        )

    def forward(self, inputs):
        z = self.encoder(inputs)
        result = self.decoder(z)
        return z, result

# ------- Load Model -------
model = torch.load('autoencoder.pth') # 載入訓練好的模型
model.eval() # 切換到測試模式

# ------- Download Testing DataSet -------
test_set = torchvision.datasets.MNIST(
    root='mnist',
    train=False, # testing data
    download=True,
    transform=torchvision.transforms.ToTensor(), # 將資料轉成tensor格式
)

# ------- DataLoader -------
test_loader = data.DataLoader(test_set, batch_size=16, shuffle=False)


# ------- Test -------
# no_grad(): 停止梯度計算
with torch.no_grad():
    for data in test_loader:
        inputs = data[0].view(-1, 28*28)
        show_images(inputs) # 顯示原始資料
        plt.show()

        z, outputs = model(inputs)
        show_images(outputs) # 顯示 AutoEncoder 生成結果
        plt.show()
        exit()

• 原始資料
  

• 生成結果

補充

• Pytorch為什麼每一輪batch都需要設置optimizer.zero_grad() ?

  根據Pytorch中backward()函數的計算，當網路參數進行反饋時，梯度是累積計算的，但在處理每一個batch時並不需要與其他batch的梯度混合起來累積計算，因此需要對每一個batch調用一遍zero_grad()將參數梯度重置為0。

  另外，如果不是在處理每個batch前都清除一次梯度，而是兩次或多次再清除一次，相當於提高了batch_size，對硬體的要求更高。

• no_grad()作用?
  停止梯度計算(停止autograd的工作)，可以加速GPU的計算和節省顯存。

• model.eval()作用?
  切換到測試模式，該模式的gradient計算和儲存跟training模式一樣，只是不會進行反向傳播(backprobagation)。