add AE example

Author: jm12138

examples/as_loss/as_loss.py

@@ -1,12 +1,14 @@
+import os
+import sys
 import paddle
-from paddle.optimizer import Adam
-from PIL import Image
 import numpy as np
-import sys, os
-import paddle.nn.functional as F
+
+from PIL import Image
+from paddle.optimizer import Adam
 sys.path.append(os.path.join(os.path.dirname(os.path.dirname(os.path.realpath(__file__))), '..'))
 from paddle_msssim import ssim, ms_ssim, SSIM, MS_SSIM
 
+
 loss_type = 'msssim'
 assert loss_type in ['ssim', 'msssim']
 

examples/auto_encoder/README.md

@@ -0,0 +1,33 @@
+# Train an autoencoder with SSIM & MS-SSIM
+
+## Prepare dataset
+* Download CLIC datase from http://clic.compression.cc/2021/tasks/index.html.
+
+* Unzip them into datasets.
+
+* The structure of the directory:
+
+    ```yaml
+    - datasets
+      - CLIC
+        - train
+          - *.png
+          - ...
+        - valid
+          - *.png
+          - ...
+    ```
+
+## Train
+* SSIM loss:
+
+    ```bash
+    $ python train.py --loss_type ssim
+    ```
+
+* MS-SSIM loss:
+
+    ```bash
+    $ python train.py --loss_type ms_ssim
+    ```
+    

examples/auto_encoder/datas/__init__.py

@@ -0,0 +1 @@
+from .image_dataset import ImageDataset

examples/auto_encoder/datas/image_dataset.py

@@ -0,0 +1,20 @@
+import os
+from PIL import Image
+from paddle.io import Dataset
+
+
+class ImageDataset(Dataset):
+    def __init__(self, root, transform=None):
+        self.root = root
+        self.transform = transform
+        self.images = list(os.listdir(root))
+        self.images.sort()
+
+    def __getitem__(self, idx):
+        img = Image.open(os.path.join(self.root, self.images[idx]))
+        if self.transform is not None:
+            img = self.transform(img)
+        return img, 
+
+    def __len__(self):
+        return len(self.images)

examples/auto_encoder/models/__init__.py

@@ -0,0 +1 @@
+from .autoencoder import AutoEncoder

examples/auto_encoder/models/autoencoder.py

@@ -0,0 +1,72 @@
+import paddle.nn as nn
+import paddle.nn.functional as F
+
+from .gdn import GDN
+
+
+# https://arxiv.org/pdf/1611.01704.pdf
+# A simplfied version without quantization
+class AutoEncoder(nn.Layer):
+    def __init__(self, C=128, M=128, in_chan=3, out_chan=3):
+        super(AutoEncoder, self).__init__()
+        self.encoder = Encoder(C=C, M=M, in_chan=in_chan)
+        self.decoder = Decoder(C=C, M=M, out_chan=out_chan)
+
+    def forward(self, x, **kargs):
+        code = self.encoder(x)
+        out = self.decoder(code)
+        return out
+
+
+class Encoder(nn.Layer):
+    """ Encoder
+    """
+
+    def __init__(self, C=32, M=128, in_chan=3):
+        super(Encoder, self).__init__()
+        self.enc = nn.Sequential(
+            nn.Conv2D(in_channels=in_chan, out_channels=M,
+                      kernel_size=5, stride=2, padding=2, bias_attr=False),
+            GDN(M),
+
+            nn.Conv2D(in_channels=M, out_channels=M, kernel_size=5,
+                      stride=2, padding=2, bias_attr=False),
+            GDN(M),
+
+            nn.Conv2D(in_channels=M, out_channels=M, kernel_size=5,
+                      stride=2, padding=2, bias_attr=False),
+            GDN(M),
+
+            nn.Conv2D(in_channels=M, out_channels=C, kernel_size=5,
+                      stride=2, padding=2, bias_attr=False)
+        )
+
+    def forward(self, x):
+        return self.enc(x)
+
+
+class Decoder(nn.Layer):
+    """ Decoder
+    """
+
+    def __init__(self, C=32, M=128, out_chan=3):
+        super(Decoder, self).__init__()
+        self.dec = nn.Sequential(
+            nn.Conv2DTranspose(in_channels=C, out_channels=M, kernel_size=5,
+                               stride=2, padding=2, output_padding=1, bias_attr=False),
+            GDN(M, inverse=True),
+
+            nn.Conv2DTranspose(in_channels=M, out_channels=M, kernel_size=5,
+                               stride=2, padding=2, output_padding=1, bias_attr=False),
+            GDN(M, inverse=True),
+
+            nn.Conv2DTranspose(in_channels=M, out_channels=M, kernel_size=5,
+                               stride=2, padding=2, output_padding=1, bias_attr=False),
+            GDN(M, inverse=True),
+
+            nn.Conv2DTranspose(in_channels=M, out_channels=out_chan, kernel_size=5,
+                               stride=2, padding=2, output_padding=1, bias_attr=False),
+        )
+
+    def forward(self, q):
+        return F.sigmoid(self.dec(q))

examples/auto_encoder/models/gdn.py

@@ -0,0 +1,47 @@
+import paddle
+import paddle.nn as nn
+import paddle.nn.functional as F
+
+
+class GDN(nn.Layer):
+    def __init__(self,
+                 num_features,
+                 inverse=False,
+                 gamma_init=.1,
+                 beta_bound=1e-6,
+                 gamma_bound=0.0,
+                 reparam_offset=2**-18,
+                 ):
+        super(GDN, self).__init__()
+        self._inverse = inverse
+        self.num_features = num_features
+        self.reparam_offset = reparam_offset
+        self.pedestal = self.reparam_offset**2
+
+        beta_init = paddle.sqrt(paddle.ones((num_features, ), dtype=paddle.float32) + self.pedestal)
+        gama_init = paddle.sqrt(paddle.full((num_features, num_features), fill_value=gamma_init, dtype=paddle.float32)
+                                * paddle.eye(num_features, dtype=paddle.float32) + self.pedestal)
+
+        self.beta = self.create_parameter(
+            shape=beta_init.shape, default_initializer=nn.initializer.Assign(beta_init))
+        self.gamma = self.create_parameter(
+            shape=gama_init.shape, default_initializer=nn.initializer.Assign(gama_init))
+
+        self.beta_bound = (beta_bound + self.pedestal) ** 0.5
+        self.gamma_bound = (gamma_bound + self.pedestal) ** 0.5
+
+    def _reparam(self, var, bound):
+        var = paddle.clip(var, min=bound)
+        return (var**2) - self.pedestal
+
+    def forward(self, x):
+        gamma = self._reparam(self.gamma, self.gamma_bound).reshape((self.num_features, self.num_features, 1, 1))  # expand to (C, C, 1, 1)
+        beta = self._reparam(self.beta, self.beta_bound)
+        norm_pool = F.conv2d(x ** 2, gamma, bias=beta, stride=1, padding=0)
+        norm_pool = paddle.sqrt(norm_pool)
+
+        if self._inverse:
+            norm_pool = x * norm_pool
+        else:
+            norm_pool = x / norm_pool
+        return norm_pool

examples/auto_encoder/train.py

@@ -0,0 +1,136 @@
+import os
+import sys
+import paddle
+import argparse
+
+from PIL import Image
+from models import AutoEncoder
+from datas import ImageDataset
+from paddle.vision import transforms
+from paddle.optimizer import Adam
+from paddle.io import DataLoader
+
+sys.path.append(os.path.join(os.path.dirname(os.path.dirname(os.path.realpath(__file__))), '..'))
+from paddle_msssim import ssim, ms_ssim, SSIM, MS_SSIM
+
+
+class MS_SSIM_Loss(MS_SSIM):
+    def forward(self, img1, img2):
+        return 100*(1 - super(MS_SSIM_Loss, self).forward(img1, img2))
+
+
+class SSIM_Loss(SSIM):
+    def forward(self, img1, img2):
+        return 100*(1 - super(SSIM_Loss, self).forward(img1, img2))
+
+
+def get_argparser():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--ckpt", default=None, type=str,
+                        help="path to trained model. Leave it None if you want to retrain your model")
+    parser.add_argument("--loss_type", type=str,
+                        default='ssim', choices=['ssim', 'ms_ssim'])
+    parser.add_argument("--batch_size", type=int, default=8)
+    parser.add_argument("--log_interval", type=int, default=10)
+    parser.add_argument("--total_epochs", type=int, default=50)
+    return parser
+
+
+def main():
+    opts = get_argparser().parse_args()
+
+    # dataset
+    train_trainsform = transforms.Compose([
+        transforms.RandomCrop(size=512, pad_if_needed=True),
+        transforms.RandomHorizontalFlip(),
+        transforms.RandomVerticalFlip(),
+        transforms.ToTensor(),
+    ])
+
+    val_transform = transforms.Compose([
+        transforms.CenterCrop(size=512),
+        transforms.ToTensor()
+    ])
+
+    train_loader = DataLoader(
+        ImageDataset(root='datasets/CLIC/train', transform=train_trainsform),
+        batch_size=opts.batch_size, shuffle=True, num_workers=0, drop_last=True)
+
+    val_loader = DataLoader(
+        ImageDataset(root='datasets/CLIC/valid', transform=val_transform),
+        batch_size=opts.batch_size, shuffle=False, num_workers=0)
+
+    print("Train set: %d, Val set: %d" %
+          (len(train_loader.dataset), len(val_loader.dataset)))
+    model = AutoEncoder(C=128, M=128, in_chan=3, out_chan=3)
+
+    # optimizer
+    optimizer = Adam(parameters=model.parameters(),
+                     learning_rate=1e-4,
+                     weight_decay=1e-5)
+
+    # checkpoint
+    best_score = 0.0
+    cur_epoch = 0
+    if opts.ckpt is not None and os.path.isfile(opts.ckpt):
+        model.set_dict(paddle.load(opts.ckpt))
+    else:
+        print("[!] Retrain")
+
+    if opts.loss_type == 'ssim':
+        criterion = SSIM_Loss(data_range=1.0, size_average=True, channel=3)
+    else:
+        criterion = MS_SSIM_Loss(data_range=1.0, size_average=True, channel=3)
+
+    #==========   Train Loop   ==========#
+    for cur_epoch in range(opts.total_epochs):
+        # =====  Train  =====
+        model.train()
+        for cur_step, (images, ) in enumerate(train_loader):
+            optimizer.clear_grad()
+            outputs = model(images)
+
+            loss = criterion(outputs, images)
+            loss.backward()
+
+            optimizer.step()
+
+            if (cur_step) % opts.log_interval == 0:
+                print("Epoch %d, Batch %d/%d, loss=%.6f" %
+                      (cur_epoch, cur_step, len(train_loader), loss.item()))
+
+        # =====  Save Latest Model  =====
+        paddle.save(model.state_dict(), 'latest_model.pdparams')
+
+        # =====  Validation  =====
+        print("Val...")
+        best_score = 0.0
+        cur_score = test(opts, model, val_loader)
+        print("%s = %.6f" % (opts.loss_type, cur_score))
+        # =====  Save Best Model  =====
+        if cur_score > best_score:  # save best model
+            best_score = cur_score
+            paddle.save(model.state_dict(), 'best_model.pdparams')
+            print("Best model saved as best_model.pt")
+
+
+def test(opts, model, val_loader):
+    model.eval()
+    cur_score = 0.0
+
+    metric = ssim if opts.loss_type == 'ssim' else ms_ssim
+
+    with paddle.no_grad():
+        for i, (images, ) in enumerate(val_loader):
+            outputs = model(images)
+            # save the first reconstructed image
+            if i == 20:
+                Image.fromarray((outputs*255).squeeze(0).detach().numpy().astype(
+                    'uint8').transpose(1, 2, 0)).save('recons_%s.png' % (opts.loss_type))
+            cur_score += metric(outputs, images, data_range=1.0)
+        cur_score /= len(val_loader.dataset)
+    return cur_score
+
+
+if __name__ == '__main__':
+    main()