pytorch自动编码器实现有损图像压缩
自动编码器（AutoEncoder）由编码器（Encoder）和解码器（Decoder）两部分组成。编码器和解码器可以是任意模型，通常神经网络模型作为编码器和解码器。
自动编码器作为一种数据压缩的方法，其原理是：输入数据经过编码器变成一个编码（code），然后将这个编码作为解码器的输入，观察解码器的输出是否能还原原始数据，因此将解码器的输出和原始数据的误差作为最优化的目标。
下面以MNIST数据集为例，使用pytorch1.0构建一个卷积神经网络做自动编码器。
1.添加引用的库文件
import os
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
from torchvision.utils import save_image2.定义超参数，是否使用GPU加速
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
batch_size = 512
3.加载MNIST数据集，并将图片的大小变为-1~1之间，这样可以使输入变得更对称，训练更加容易收敛。
# 标准化
data_tf = transforms.Compose(
[transforms.ToTensor(),
transforms.Normalize([0.5], [0.5])]
)
train_dataset = datasets.MNIST(root='./data', train=True, transform=data_tf, download=True)
train_data = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
4.定义卷积神经网络的自动编码器
class AutoEncoder(nn.Module):
def __init__(self):
super(AutoEncoder, self).__init__()
self.encoder = nn.Sequential(
nn.Conv2d(1, 16, 3, stride=3, padding=1), # b,16,10,10
nn.ReLU(True),
nn.MaxPool2d(2, stride=2), # b,16,5,5
nn.Conv2d(16, 8, 3, stride=2, padding=1), # b,8,3,3
nn.ReLU(True),
nn.MaxPool2d(2, stride=1) # b,8,2,2
)
self.decoder = nn.Sequential(
nn.ConvTranspose2d(8, 16, 3, stride=2), # b,16,5,5
nn.ReLU(True),
nn.ConvTranspose2d(16, 8, 5, stride=3, padding=1), # b,8,15,15
nn.ReLU(True),
nn.ConvTranspose2d(8, 1, 2, stride=2, padding=1), # b,1,28,28
nn.Tanh()
)
def forward(self, x):
encode = self.encoder(x)
decode = self.decoder(encode)
return encode, decode
torch.nn.Conv2d(in_channels, out_channels, kernel_size, stride=1, padding=0,groups=1, bias=True, dilation=1)
- in_channels(int):输入数据的通道数；
- out_channels(int)：输出数据的通道数；
- kernel_size(int or tuple)：滤波器或卷积核的大小；
- stride(int or tuple,optional) ：步长；
- padding(int or tuple, optional)：四周是否进行0填充；
- groups(int, optional) – 从输入通道到输出通道的阻塞连接数
- bias(bool, optional) - 如果bias=True，添加偏置
- dilation(int or tuple, optional) – 卷积核元素之间的间距
对于每一条边输入,输出的尺寸的公式如下：
解码器使用nn.ConvTranspose2d()，可以看作卷积的反操作。具体参数如下：
torch.nn.ConvTranspose2d(in_channels, out_channels, kernel_size, stride=1, padding=0,output_padding=0, groups=1, bias=True, dilation=1)
- in_channels(int) – 输入信号的通道数
- out_channels(int) – 卷积产生的通道数
- kerner_size(int or tuple) - 卷积核的大小
- stride(int or tuple,optional) - 卷积步长，即要将输入扩大的倍数。
- padding(int or tuple, optional) - 输入的每一条边补充0的层数，高宽都增加2*padding
- output_padding(int or tuple, optional) - 输出边补充0的层数，高宽都增加padding
- groups(int, optional) – 从输入通道到输出通道的阻塞连接数
- bias(bool, optional) - 如果bias=True，添加偏置
- dilation(int or tuple, optional) – 卷积核元素之间的间距
对于每一条边输入,输出的尺寸的公式如下：
5.实例化模型，定义loss函数和优化函数
model = AutoEncoder().to(device)
# 定义loss函数和优化方法
loss_fn = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=1e-3, weight_decay=1e-5)6.训练并保存解码器生成的图片
for t in range(40):
for data in train_data:
img, label = data
img = img.to(device)
label = label.to(device)
_, output = model(img)
loss = loss_fn(output, img) / img.shape[0] # 平均损失
# 反向传播
optimizer.zero_grad()
loss.backward()
optimizer.step()
if (t + 1) % 5 == 0: # 每 5 次，保存一下解码的图片和原图片
print('epoch: {}, Loss: {:.4f}'.format(t + 1, loss.item()))
pic = to_img(output.cpu().data)
if not os.path.exists('./conv_autoencoder'):
os.mkdir('./conv_autoencoder')
save_image(pic, './conv_autoencoder/decode_image_{}.png'.format(t + 1))
save_image(img, './conv_autoencoder/raw_image_{}.png'.format(t + 1))结果对比（左边生成图片，右边原始图片）：
附上完整代码：
import os
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
from torchvision.utils import save_image
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
batch_size = 512
# 标准化
data_tf = transforms.Compose(
[transforms.ToTensor(),
transforms.Normalize([0.5], [0.5])]
)
train_dataset = datasets.MNIST(root='./data', train=True, transform=data_tf, download=True)
train_data = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
def to_img(x):
x = 0.5 * (x + 1.) # 将-1~1转成0-1
x = x.clamp(0, 1)
x = x.view(x.shape[0], 1, 28, 28)
return x
class AutoEncoder(nn.Module):
def __init__(self):
super(AutoEncoder, self).__init__()
self.encoder = nn.Sequential(
nn.Conv2d(1, 16, 3, stride=3, padding=1), # b,16,10,10
nn.ReLU(True),
nn.MaxPool2d(2, stride=2), # b,16,5,5
nn.Conv2d(16, 8, 3, stride=2, padding=1), # b,8,3,3
nn.ReLU(True),
nn.MaxPool2d(2, stride=1) # b,8,2,2
)
self.decoder = nn.Sequential(
# nn.ConvTranspose2d(8, 8, 3, stride=2, padding=1), # b,8,3,3
# nn.ReLU(True),
# nn.ConvTranspose2d(8, 16, 4, stride=4, padding=1), # b,16,10,10
# nn.ReLU(True),
# nn.ConvTranspose2d(16, 1, 3, stride=3, padding=1), # b,1,28,28
# nn.Tanh()
nn.ConvTranspose2d(8, 16, 3, stride=2), # b,16,5,5
nn.ReLU(True),
nn.ConvTranspose2d(16, 8, 5, stride=3, padding=1), # b,8,15,15
nn.ReLU(True),
nn.ConvTranspose2d(8, 1, 2, stride=2, padding=1), # b,1,28,28
nn.Tanh()
)
def forward(self, x):
encode = self.encoder(x)
decode = self.decoder(encode)
return encode, decode
model = AutoEncoder().to(device)
# 定义loss函数和优化方法
loss_fn = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=1e-3, weight_decay=1e-5)
for t in range(40):
for data in train_data:
img, label = data
img = img.to(device)
label = label.to(device)
_, output = model(img)
loss = loss_fn(output, img) / img.shape[0] # 平均损失
# 反向传播
optimizer.zero_grad()
loss.backward()
optimizer.step()
if (t + 1) % 5 == 0: # 每 5 次，保存一下解码的图片和原图片
print('epoch: {}, Loss: {:.4f}'.format(t + 1, loss.item()))
pic = to_img(output.cpu().data)
if not os.path.exists('./conv_autoencoder'):
os.mkdir('./conv_autoencoder')
save_image(pic, './conv_autoencoder/decode_image_{}.png'.format(t + 1))
save_image(img, './conv_autoencoder/raw_image_{}.png'.format(t + 1))