Add files

Author: moelife-coder

Test.py

@@ -0,0 +1,34 @@
+import argparse
+import torchvision
+import torch
+from PIL import Image
+
+
+def test(data, network):
+    output, _ = network(data)
+    return output
+
+
+parser = argparse.ArgumentParser(
+    description = 'Testing utility for paper "Investigation on different loss function in image autoencoder"')
+parser.add_argument('target', type = str, help = "Input image")
+parser.add_argument('output', type = str, help = "Output image")
+parser.add_argument('--model_name', type = str, dest = "model_name", help = "Model name for saving & restoring model",
+                    default = "saved.pkl")
+parser.add_argument('--disable-cuda', action='store_true',
+                    help='Disable CUDA')
+args = parser.parse_args()
+filename = args.target
+output = args.output
+model_name = args.model_name
+network = torch.load(model_name)
+if not args.disable_cuda and torch.cuda.is_available():
+    args.device = torch.device('cuda')
+    network.cuda()
+else:
+    args.device = torch.device('cpu')
+data = Image.open(filename)
+data = torchvision.transforms.ToTensor()(data).unsqueeze_(0).to(args.device)
+result = test(data, network)
+result = torchvision.transforms.ToPILImage()(result.squeeze_(0).to("cpu"))
+result.save(output)

Train.py

@@ -0,0 +1,109 @@
+import argparse
+import torchvision
+import torch
+from torch.utils import data
+
+
+def train(EPOCH, dataloader, optimizer, loss_function, network, model_name, negative_loss, device, should_view, f):
+    for epoch in range(EPOCH):
+        loss = 0
+        for step, (b_x, _) in enumerate(dataloader):
+            b_x = b_x.to(device)
+            if should_view:
+                formatted_b_x = b_x.view(b_x.shape[0], -1)
+            else:
+                formatted_b_x = b_x
+            output, _ = network(formatted_b_x)
+            output = output.view(b_x.shape)
+            loss = loss_function(output, b_x)
+            if negative_loss:
+                loss = - loss
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+            print('Epoch: ', epoch + 1, '| Step: ', step + 1, '| Train loss: %.4f' % loss.data.numpy())
+            if step % SAVE_STEP == 1:
+                print('Saving models...')
+                torch.save(network, model_name)
+        print('Saving models...')
+        torch.save(network, model_name)
+        print('Saving logs...')
+        f.write("{} {}\n".format(epoch, loss.data.numpy()))
+
+
+parser = argparse.ArgumentParser(
+    description = 'Training utility for paper "Investigation on different loss function in image autoencoder"')
+parser.add_argument('--save_at', type = int, dest = "SAVE_STEP", help = "Save network at how many steps (default: 10)  "
+                                                                        ",", default = 10)
+parser.add_argument('--epoch', '-e', type = int, dest = "EPOCH", help = "Epoch for training", default = 10)
+parser.add_argument('--batch_size', '-b', type = int, dest = "BATCH_SIZE", help = "Batch size for dataloader",
+                    default = 8192)
+parser.add_argument('--learning_rate', '-l', type = float, dest = "LR", help = "Learning rate", default = 0.001)
+parser.add_argument('--dataset', type = str, dest = "dataset", default = "cifar10", choices = ['cifar10', 'mnist'])
+parser.add_argument('--network', type = str, dest = "network", default = "mlp", choices = ['mlp', 'conv'])
+parser.add_argument('origional_size', type = int, help = "Size of origional image. Please note that for ConvNet, it is the number of images' channel.")
+parser.add_argument('bottleneck', type = int, help = "Size of bottle neck. Please note that for ConvNet, its \"bottleneck\" is input * origional size in one channel / 4")
+parser.add_argument('--loss_func', type = str, dest = "loss_function", default = "mse",
+                    choices = ['mse', 'l1', 'ssim', 'psnr'])
+parser.add_argument('--model_name', type = str, dest = "model_name", help = "Model name for saving & restoring model",
+                    default = "saved.pkl")
+parser.add_argument('--disable-cuda', action='store_true',
+                    help='Disable CUDA')
+parser.add_argument('--log', type = str, default = "train.log",dest = "log_file", help = "Plase to store logs")
+args = parser.parse_args()
+SAVE_STEP = args.SAVE_STEP
+EPOCH = args.EPOCH
+BATCH_SIZE = args.BATCH_SIZE
+LR = args.LR
+dataset = args.dataset
+network = args.network
+orgsize = args.origional_size
+bottleneck = args.bottleneck
+loss_function = args.loss_function
+model_name = args.model_name
+log_file = args.log_file
+f = open(log_file, "a")
+f.write("x y\n")
+if dataset == "cifar10":
+    train_data = torchvision.datasets.CIFAR10(
+        root = './cifar10/',
+        transform = torchvision.transforms.ToTensor(),
+        download = True,
+    )
+else:
+    train_data = torchvision.datasets.MNIST(
+        root = './mnist/',
+        transform = torchvision.transforms.ToTensor(),
+        download = True,
+    )
+if network == "mlp":
+    import mlp_network
+    should_view = True
+    network = mlp_network.autoencoder(orgsize, bottleneck)
+else:
+    import conv_network
+    should_view = False
+    network = conv_network.autoencoder(orgsize, bottleneck)
+if not args.disable_cuda and torch.cuda.is_available():
+    args.device = torch.device('cuda')
+    network.cuda()
+else:
+    args.device = torch.device('cpu')
+negative_loss = False
+if loss_function == "mse":
+    loss_func = torch.nn.MSELoss()
+elif loss_function == "l1":
+    loss_func = torch.nn.L1Loss()
+elif loss_function == "l1s":
+    loss_func = torch.nn.SmoothL1Loss()
+elif loss_function == "ssim":
+    import pytorch_ssim
+    loss_func = pytorch_ssim.SSIM()
+    negative_loss = True
+else:
+    import psnr
+    loss_func = psnr.PSNR()
+train_loader = data.DataLoader(dataset = train_data, batch_size = BATCH_SIZE, shuffle = True)
+optimizer = torch.optim.Adam(network.parameters(), lr = LR)
+train(EPOCH, train_loader, optimizer, loss_func, network, model_name, negative_loss, args.device, should_view, f)
+f.close()

conv_network.py

@@ -0,0 +1,65 @@
+from torch import nn
+
+
+class encoder(nn.Module):
+    def __init__(self, in_features, out_features):
+        super(encoder, self).__init__()
+        self.conv1 = nn.Conv2d(
+            in_channels = in_features,
+            out_channels = 32,
+            kernel_size = 3,
+            padding = 1
+        )
+        self.conv2 = nn.Conv2d(
+            in_channels = 32,
+            out_channels = 64,
+            kernel_size = 3,
+            padding = 1
+        )
+        self.pool = nn.MaxPool2d(2, 0)
+        self.conv3 = nn.Conv2d(
+            in_channels = 64,
+            out_channels = out_features,
+            kernel_size = 3,
+            padding = 1
+        )
+
+    def forward(self, x):
+        return self.conv3(self.pool(self.conv2(self.conv1(x))))
+
+class decoder(nn.Module):
+    def __init__(self, in_features, out_features):
+        super(decoder, self).__init__()
+        self.conv1 = nn.Conv2d(
+            in_channels = in_features,
+            out_channels = 64,
+            kernel_size = 3,
+            padding = 1
+        )
+        self.pool = nn.Upsample(scale_factor = 2)
+        self.conv2 = nn.Conv2d(
+            in_channels = 64,
+            out_channels = 32,
+            kernel_size = 3,
+            padding = 1
+        )
+        self.conv3 = nn.Conv2d(
+            in_channels = 32,
+            out_channels = out_features,
+            kernel_size = 3,
+            padding = 1
+        )
+
+    def forward(self, x):
+        return self.conv3(self.conv2(self.pool(self.conv1(x))))
+
+
+class autoencoder(nn.Module):
+    def __init__(self, in_feature, bottleneck):
+        super(autoencoder, self).__init__()
+        self.encoder = encoder(in_feature, bottleneck)
+        self.decoder = decoder(bottleneck, in_feature)
+
+    def forward(self, x):
+        encoded = self.encoder(x)
+        return self.decoder(encoded), encoded

mlp_network.py

@@ -0,0 +1,21 @@
+from torch import nn
+
+
+class coder(nn.Module):
+    def __init__(self, in_features, out_features):
+        super(coder, self).__init__()
+        self.linear1 = nn.Linear(in_features, out_features)
+
+    def forward(self, x):
+        return self.linear1(x)
+
+
+class autoencoder(nn.Module):
+    def __init__(self, in_feature, bottleneck):
+        super(autoencoder, self).__init__()
+        self.encoder = coder(in_feature, bottleneck)
+        self.decoder = coder(bottleneck, in_feature)
+
+    def forward(self, x):
+        encoded = self.encoder(x)
+        return self.decoder(encoded), encoded

psnr.py

@@ -0,0 +1,12 @@
+import torch
+from math import log10
+
+
+class PSNR(torch.nn.Module):
+    def __init__(self):
+        super(PSNR, self).__init__()
+        self.criterion = torch.nn.MSELoss()
+    def forward(self, prediction, target):
+        mse = self.criterion(prediction, target)
+        psnr = 10 * log10(1 / mse.item())
+        return psnr

pytorch_ssim/__init__.py

@@ -0,0 +1,73 @@
+import torch
+import torch.nn.functional as F
+from torch.autograd import Variable
+import numpy as np
+from math import exp
+
+def gaussian(window_size, sigma):
+    gauss = torch.Tensor([exp(-(x - window_size//2)**2/float(2*sigma**2)) for x in range(window_size)])
+    return gauss/gauss.sum()
+
+def create_window(window_size, channel):
+    _1D_window = gaussian(window_size, 1.5).unsqueeze(1)
+    _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0)
+    window = Variable(_2D_window.expand(channel, 1, window_size, window_size).contiguous())
+    return window
+
+def _ssim(img1, img2, window, window_size, channel, size_average = True):
+    mu1 = F.conv2d(img1, window, padding = window_size//2, groups = channel)
+    mu2 = F.conv2d(img2, window, padding = window_size//2, groups = channel)
+
+    mu1_sq = mu1.pow(2)
+    mu2_sq = mu2.pow(2)
+    mu1_mu2 = mu1*mu2
+
+    sigma1_sq = F.conv2d(img1*img1, window, padding = window_size//2, groups = channel) - mu1_sq
+    sigma2_sq = F.conv2d(img2*img2, window, padding = window_size//2, groups = channel) - mu2_sq
+    sigma12 = F.conv2d(img1*img2, window, padding = window_size//2, groups = channel) - mu1_mu2
+
+    C1 = 0.01**2
+    C2 = 0.03**2
+
+    ssim_map = ((2*mu1_mu2 + C1)*(2*sigma12 + C2))/((mu1_sq + mu2_sq + C1)*(sigma1_sq + sigma2_sq + C2))
+
+    if size_average:
+        return ssim_map.mean()
+    else:
+        return ssim_map.mean(1).mean(1).mean(1)
+
+class SSIM(torch.nn.Module):
+    def __init__(self, window_size = 11, size_average = True):
+        super(SSIM, self).__init__()
+        self.window_size = window_size
+        self.size_average = size_average
+        self.channel = 1
+        self.window = create_window(window_size, self.channel)
+
+    def forward(self, img1, img2):
+        (_, channel, _, _) = img1.size()
+
+        if channel == self.channel and self.window.data.type() == img1.data.type():
+            window = self.window
+        else:
+            window = create_window(self.window_size, channel)
+            
+            if img1.is_cuda:
+                window = window.cuda(img1.get_device())
+            window = window.type_as(img1)
+            
+            self.window = window
+            self.channel = channel
+
+
+        return _ssim(img1, img2, window, self.window_size, channel, self.size_average)
+
+def ssim(img1, img2, window_size = 11, size_average = True):
+    (_, channel, _, _) = img1.size()
+    window = create_window(window_size, channel)
+    
+    if img1.is_cuda:
+        window = window.cuda(img1.get_device())
+    window = window.type_as(img1)
+    
+    return _ssim(img1, img2, window, window_size, channel, size_average)

requirements.txt

@@ -0,0 +1,5 @@
+torchvision>=0.6.0
+torch>=1.5.0
+Pillow>=7.1.2
+setuptools>=46.4.0
+numpy>=1.18.4

run.sh

@@ -0,0 +1,16 @@
+#!/bin/bash
+test_img() {
+  python Test.py --model_name "$1" "$2" "$3"
+}
+train_mlp_mnist() {
+  python Train.py --epoch 20 --dataset mnist --network mlp --loss_func "$2" --model_name "$3" 784 "$1"
+}
+train_cnn_mnist() {
+  python Train.py --epoch 20 --dataset mnist --network cnn --loss_func "$2" --model_name "$3" 1 "$1"
+}
+train_mlp_cifar() {
+  python Train.py --epoch 20 --dataset cifar10 --network mlp --loss_func "$2" --model_name "$3" 784 "$1"
+}
+train_cnn_cifar() {
+  python Train.py --epoch 20 --dataset cifar10 --network cnn --loss_func "$2" --model_name "$3" 1 "$1"
+}

train.log

@@ -0,0 +1,22 @@
+x y
+x y
+x y
+x y
+x y
+x y
+x y
+x y
+x y
+x y
+x y
+x y
+x y
+x y
+x y
+x y
+x y
+x y
+x y
+x y
+x y
+0 0.14592687785625458