Add files via upload

Author: YangiD

DBCNN/DBCNN_train_attack.py

@@ -0,0 +1,289 @@
+import os
+
+import torch
+import torchvision
+import torch.nn.functional as F
+import torch.nn as nn
+import numpy as np
+import DBCNN.WPFolder
+from PIL import Image
+
+torch.manual_seed(0)
+torch.cuda.manual_seed_all(0)
+
+
+def pil_loader(path):
+    # open path as file to avoid ResourceWarning (https://github.com/python-pillow/Pillow/issues/835)
+    with open(path, 'rb') as f:
+        img = Image.open(f)
+        return img.convert('RGB')
+
+
+def accimage_loader(path):
+    import accimage
+    try:
+        return accimage.Image(path)
+    except IOError:
+        # Potentially a decoding problem, fall back to PIL.Image
+        return pil_loader(path)
+
+
+def default_loader(path):
+    from torchvision import get_image_backend
+    if get_image_backend() == 'accimage':
+        return accimage_loader(path)
+    else:
+        return pil_loader(path)
+    
+IMG_EXTENSIONS = ['.jpg', '.jpeg', '.png', '.ppm', '.bmp', '.pgm', '.tif']
+
+def weight_init(net): 
+    for m in net.modules():    
+        if isinstance(m, nn.Conv2d):         
+            nn.init.kaiming_normal_(m.weight.data,nonlinearity='relu')
+            m.bias.data.zero_()
+        elif isinstance(m, nn.Linear):
+            nn.init.kaiming_normal_(m.weight.data,nonlinearity='relu')
+            m.bias.data.zero_()
+        elif isinstance(m, nn.BatchNorm2d):
+            m.weight.data.fill_(1)
+            m.bias.data.zero_()
+        
+
+
+class SCNN(nn.Module):
+
+    def __init__(self):
+        """Declare all needed layers."""
+        super(SCNN, self).__init__()
+
+        # Linear classifier.
+
+        self.num_class = 39
+#        self.features = nn.Sequential(nn.Conv2d(3,48,3,1,1),nn.ReLU(inplace=True),
+#                                      nn.Conv2d(48,48,3,2,1),nn.ReLU(inplace=True),
+#                                      nn.Conv2d(48,64,3,1,1),nn.ReLU(inplace=True),
+#                                      nn.Conv2d(64,64,3,2,1),nn.ReLU(inplace=True),
+#                                      nn.Conv2d(64,64,3,1,1),nn.ReLU(inplace=True),
+#                                      nn.Conv2d(64,64,3,2,1),nn.ReLU(inplace=True),
+#                                      nn.Conv2d(64,128,3,1,1),nn.ReLU(inplace=True),
+#                                      nn.Conv2d(128,128,3,1,1),nn.ReLU(inplace=True),
+#                                      nn.Conv2d(128,128,3,2,1),nn.ReLU(inplace=True))
+        self.features = nn.Sequential(nn.Conv2d(3,48,3,1,1),nn.BatchNorm2d(48),nn.ReLU(inplace=True),
+                                      nn.Conv2d(48,48,3,2,1),nn.BatchNorm2d(48),nn.ReLU(inplace=True),
+                                      nn.Conv2d(48,64,3,1,1),nn.BatchNorm2d(64),nn.ReLU(inplace=True),
+                                      nn.Conv2d(64,64,3,2,1),nn.BatchNorm2d(64),nn.ReLU(inplace=True),
+                                      nn.Conv2d(64,64,3,1,1),nn.BatchNorm2d(64),nn.ReLU(inplace=True),
+                                      nn.Conv2d(64,64,3,2,1),nn.BatchNorm2d(64),nn.ReLU(inplace=True),
+                                      nn.Conv2d(64,128,3,1,1),nn.BatchNorm2d(128),nn.ReLU(inplace=True),
+                                      nn.Conv2d(128,128,3,1,1),nn.BatchNorm2d(128),nn.ReLU(inplace=True),
+                                      nn.Conv2d(128,128,3,2,1),nn.BatchNorm2d(128),nn.ReLU(inplace=True))
+        weight_init(self.features)
+        self.pooling = nn.AvgPool2d(14,1)
+        self.projection = nn.Sequential(nn.Conv2d(128,256,1,1,0), nn.BatchNorm2d(256), nn.ReLU(inplace=True),
+                                        nn.Conv2d(256,256,1,1,0), nn.BatchNorm2d(256), nn.ReLU(inplace=True))
+        weight_init(self.projection)    
+        self.classifier = nn.Linear(256,self.num_class)
+        weight_init(self.classifier)
+
+    def forward(self, X):
+#        return X
+        N = X.size()[0]
+        assert X.size() == (N, 3, 224, 224)
+        X = self.features(X)
+        assert X.size() == (N, 128, 14, 14)
+        X = self.pooling(X)
+        assert X.size() == (N, 128, 1, 1)
+        X = self.projection(X)
+        X = X.view(X.size(0), -1)          
+        X = self.classifier(X)
+        assert X.size() == (N, self.num_class)
+        return X
+
+class SCNNManager(object):
+    """Manager class to train S-CNN.
+    """
+    def __init__(self, options, path):
+        """Prepare the network, criterion, solver, and data.
+        Args:
+            options, dict: Hyperparameters.
+        """
+        print('Prepare the network and data.')
+        self._options = options
+        self._path = path
+        self._epoch = 0
+        # Network.
+        network = SCNN()
+        weight_init(network)
+        #self._net = network.cuda()
+        self._net = torch.nn.DataParallel(network).cuda()
+        
+        logspaced_LR = np.logspace(-1,-4, self._options['epochs'])   
+        # Load the model from disk.
+        checkpoints_list = os.listdir(self._path['model'])
+        if len(checkpoints_list) != 0:
+            self._net.load_state_dict(torch.load(os.path.join(self._path['model'],'%s%s%s' % ('net_params', str(len(checkpoints_list)-1), '.pkl'))))
+            self._epoch = len(checkpoints_list)
+            self._options['base_lr'] = logspaced_LR[len(checkpoints_list)]
+        #self._net.load_state_dict(torch.load(self._path['model']))
+        print(self._net)
+        # Criterion.  
+        self._criterion = torch.nn.CrossEntropyLoss().cuda()
+        # Solver.
+        self._solver = torch.optim.SGD(
+            self._net.parameters(), lr=self._options['base_lr'],
+            momentum=0.9, weight_decay=self._options['weight_decay'])
+#        self._solver = torch.optim.Adam(
+#            self._net.parameters(), lr=self._options['base_lr'],
+#            weight_decay=self._options['weight_decay'])        
+        lambda1 = lambda epoch: logspaced_LR[epoch]
+        self._scheduler = torch.optim.lr_scheduler.LambdaLR(self._solver,lr_lambda=lambda1)
+
+        train_transforms = torchvision.transforms.Compose([
+            torchvision.transforms.Resize(size=256),  # Let smaller edge match
+            torchvision.transforms.RandomHorizontalFlip(),
+            torchvision.transforms.RandomCrop(size=224),
+            torchvision.transforms.ToTensor(),
+            torchvision.transforms.Normalize(mean=(0.485, 0.456, 0.406),
+                                             std=(0.229, 0.224, 0.225))
+        ])
+        test_transforms = torchvision.transforms.Compose([
+            torchvision.transforms.Resize(size=256),
+            torchvision.transforms.CenterCrop(size=224),
+            torchvision.transforms.ToTensor(),
+            torchvision.transforms.Normalize(mean=(0.485, 0.456, 0.406),
+                                             std=(0.229, 0.224, 0.225))
+        ])
+        train_data = WPFolder.WPFolder(
+            root=self._path['waterloo_pascal'], loader = default_loader, extensions = IMG_EXTENSIONS,
+            transform=train_transforms,train = True, ratio = 0.8)
+        test_data = WPFolder.WPFolder(
+            root=self._path['waterloo_pascal'], loader = default_loader, extensions = IMG_EXTENSIONS,
+            transform=test_transforms, train = False, ratio = 0.8)
+        self._train_loader = torch.utils.data.DataLoader(
+            train_data, batch_size=self._options['batch_size'],
+            shuffle=True, num_workers=0, pin_memory=True)
+        self._test_loader = torch.utils.data.DataLoader(
+            test_data, batch_size=self._options['batch_size'],
+            shuffle=False, num_workers=0, pin_memory=True)
+
+    def train(self):
+        """Train the network."""
+        print('Training.')
+        best_acc = 0.0
+        best_epoch = None
+        print('Epoch\tTrain loss\tTrain acc\tTest acc')
+        for t in range(self._epoch,self._options['epochs']):
+            epoch_loss = []
+            num_correct = 0.0
+            num_total = 0.0
+            batchindex = 0
+            for X, y in self._train_loader:
+               X = torch.tensor(X.cuda())
+               y = torch.tensor(y.cuda(non_blocking=True)) #async=True
+               #y = torch.tensor(y.to(device))
+
+               # Clear the existing gradients.
+               self._solver.zero_grad()
+               # Forward pass.
+               score = self._net(X)
+               loss = self._criterion(score, y.detach())     
+               epoch_loss.append(loss.item())
+  
+               # Prediction.
+               _, prediction = torch.max(F.softmax(score.data), 1)
+               num_total += y.size(0)
+               num_correct += torch.sum(prediction == y)
+               # Backward pass.
+               loss.backward()
+               self._solver.step()
+               batchindex = batchindex + 1
+            print('%d epoch done' % (t+1))
+            train_acc = 100 * num_correct.float() / num_total               
+            if (t < 2) | (t > 20):
+                with torch.no_grad():
+                    test_acc = self._accuracy(self._test_loader)   
+                if test_acc > best_acc:
+                    best_acc = test_acc
+                    best_epoch = t + 1
+            print('*', end='')
+            print('%d\t%4.3f\t\t%4.2f%%\t\t%4.2f%%' %
+                  (t+1, sum(epoch_loss) / len(epoch_loss), train_acc, test_acc)) 
+            pwd = os.getcwd()
+            modelpath = os.path.join(pwd,'models',('net_params' + str(t) + '.pkl'))
+            torch.save(self._net.state_dict(), modelpath)
+            self._scheduler.step(t)
+        print('Best at epoch %d, test accuaray %f' % (best_epoch, best_acc))
+
+    def _accuracy(self, data_loader):
+        """Compute the train/test accuracy.
+        Args:
+            data_loader: Train/Test DataLoader.
+        Returns:
+            Train/Test accuracy in percentage.
+        """
+        self._net.eval()
+        num_correct = 0.0
+        num_total = 0.0
+        batchindex = 0
+        for X, y in data_loader:
+            # Data.
+            batchindex = batchindex + 1
+            X = torch.tensor(X.cuda())
+            y = torch.tensor(y.cuda(non_blocking=True)) #async=True
+
+            # Prediction.
+            score = self._net(X)
+            _, prediction = torch.max(score.data, 1)
+            num_total += y.size(0)
+            num_correct += torch.sum(prediction == y.data)
+        self._net.train()  # Set the model to training phase
+        return 100 * num_correct.float() / num_total
+
+
+
+def main():
+    """The main function."""
+    import argparse
+    parser = argparse.ArgumentParser(
+        description='Train DB-CNN for BIQA.')
+    parser.add_argument('--base_lr', dest='base_lr', type=float, default=1e-1,
+                        help='Base learning rate for training.')
+    parser.add_argument('--batch_size', dest='batch_size', type=int,
+                        default=128, help='Batch size.')
+    parser.add_argument('--epochs', dest='epochs', type=int,
+                        default=30, help='Epochs for training.')
+    parser.add_argument('--weight_decay', dest='weight_decay', type=float,
+                        default=5e-4, help='Weight decay.')
+
+    args = parser.parse_args()
+    if args.base_lr <= 0:
+        raise AttributeError('--base_lr parameter must >0.')
+    if args.batch_size <= 0:
+        raise AttributeError('--batch_size parameter must >0.')
+    if args.epochs < 0:
+        raise AttributeError('--epochs parameter must >=0.')
+    if args.weight_decay <= 0:
+        raise AttributeError('--weight_decay parameter must >0.')
+
+
+    options = {
+        'base_lr': args.base_lr,
+        'batch_size': args.batch_size,
+        'epochs': args.epochs,
+        'weight_decay': args.weight_decay,
+    }
+
+
+    path = {
+        'waterloo_pascal': 'Z:\Waterloo\exploration_database_and_code\image',
+        'model': 'D:\zwx_Project\dbcnn_pytorch\models'
+    }
+
+    manager = SCNNManager(options, path)
+    # manager.getStat()
+    manager.train()
+
+
+if __name__ == '__main__':
+    main()