Initial commit

Author: DenisDsh

compute_mean_std.py

@@ -0,0 +1,59 @@
+import torch
+from torchvision import transforms, datasets
+import numpy as np
+
+
+def compute_mean_std(path_dataset):
+    """
+    Compute mean and standard deviation of an image dataset.
+    Acknowledgment : http://forums.fast.ai/t/image-normalization-in-pytorch/7534
+    """
+    transform = transforms.Compose([
+        transforms.Resize(224),
+        transforms.ToTensor()
+    ])
+
+    dataset = datasets.ImageFolder(root=path_dataset,
+                                   transform=transform)
+    # Choose a large batch size to better approximate. Optimally load the dataset entirely on memory.
+    data_loader = torch.utils.data.DataLoader(dataset, batch_size=4096, shuffle=False, num_workers=4)
+
+    pop_mean = []
+    pop_std = []
+
+    for i, data in enumerate(data_loader, 0):
+        # shape (batch_size, 3, height, width)
+        numpy_image = data[0].numpy()
+
+        # shape (3,) -> 3 channels
+        batch_mean = np.mean(numpy_image, axis=(0, 2, 3))
+        batch_std = np.std(numpy_image, axis=(0, 2, 3))
+
+        pop_mean.append(batch_mean)
+        pop_std.append(batch_std)
+
+    # shape (num_iterations, 3) -> (mean across 0th axis) -> shape (3,)
+    pop_mean = np.array(pop_mean).mean(axis=0)
+    pop_std = np.array(pop_std).mean(axis=0)
+
+    values = {
+        'mean': pop_mean,
+        'std': pop_std
+    }
+
+    return values
+
+
+def main():
+    mean_std = {}
+    for dataset in ['amazon', 'dslr', 'webcam']:
+        # Construct path
+        dataset_path = './data/%s/images' % dataset
+        values = compute_mean_std(dataset_path)
+        # Add values to dict
+        mean_std[dataset] = values
+
+    print(mean_std)
+
+if __name__ == '__main__':
+    main()

config.py

@@ -0,0 +1,10 @@
+# Paper: In the training phase, we set the batch size to 128,
+# base learning rate to 10−3, weight decay to 5×10−4, and momentum to 0.9
+
+lr = 1e-3
+decay = 5e-4
+momentum = 0.9
+batch_size = 128
+epochs = 20
+n_classes = 31
+lambda_coral = 0

coral.py

@@ -0,0 +1,31 @@
+import torch
+import numpy as np
+
+
+def coral(source, target):
+
+    d = source.size(1)  # dim vector
+
+    source_c = compute_covariance(source)
+    target_c = compute_covariance(target)
+
+    loss = torch.sum(torch.mul((source_c - target_c), (source_c - target_c)))
+
+    loss = loss / (4 * d * d)
+    return loss
+
+
+def compute_covariance(input_data):
+    """
+    Compute Covariance matrix of the input data
+    """
+    n = input_data.size(0)  # batch_size
+
+    id_row = torch.ones(n).resize(1, n)
+    sum_column = torch.mm(id_row, input_data)
+    mean_column = torch.div(sum_column, n)
+    term_mul_2 = torch.mm(mean_column.t(), mean_column)
+    d_t_d = torch.mm(input_data.t(), input_data)
+    c = torch.add(d_t_d, (-1 * term_mul_2)) * 1 / (n - 1)
+
+    return c

data_loader.py

@@ -0,0 +1,38 @@
+import torch
+from torchvision import transforms, datasets
+import numpy as np
+import matplotlib.pyplot as plt
+
+
+def get_loader(name_dataset, batch_size, train=True):
+
+    # Computed with compute_mean_std.py
+    mean_std = {
+        'amazon': {
+            'mean': [0.79235494, 0.7862071 , 0.78418255],
+            'std':  [0.31496558, 0.3174693 , 0.3193569 ]
+        },
+        'dslr': {
+            'mean': [0.47086468, 0.44865608, 0.40637794],
+            'std':  [0.20395322, 0.19204104, 0.1996422 ]
+        },
+        'webcam': {
+            'mean': [0.6119875 , 0.6187739 , 0.61730677],
+            'std':  [0.25063968, 0.25554898, 0.25773206]
+        }
+    }
+
+    data_transform = transforms.Compose([
+            transforms.Scale(224),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=mean_std[name_dataset]['mean'],
+                                 std=mean_std[name_dataset]['std'])
+        ])
+
+    dataset = datasets.ImageFolder(root='./data/%s/images' % name_dataset,
+                                   transform=data_transform)
+    dataset_loader = torch.utils.data.DataLoader(dataset,
+                                                 batch_size=batch_size, shuffle=train,
+                                                 num_workers=4)
+    return dataset_loader
+

loading.ipynb

@@ -0,0 +1,26 @@
+from torchvision.models import alexnet
+import torch.nn as nn
+
+
+class Net(nn.Module):
+
+    def __init__(self, num_classes, pretrained=False):
+        super(Net, self).__init__()
+
+        # check https://github.com/pytorch/vision/blob/master/torchvision/models/alexnet.py
+        self.model = alexnet(pretrained=pretrained, num_classes=num_classes)
+
+        # if we want to feed 448x448 images
+        # self.model.avgpool = nn.AdaptiveAvgPool2d(1)
+
+        # In case we want to apply the loss to any other layer than the last
+        # we need a forward hook on that layer
+        # def save_features_layer_x(module, input, output):
+        #     self.layer_x = output
+
+        # This is a forward hook. Is executed each time forward is executed
+        # self.model.layer4.register_forward_hook(save_features_layer_x)
+
+    def forward(self, x):
+        out = self.model(x)
+        return out  # , self.layer_x

model.py

@@ -0,0 +1,117 @@
+from torchvision.models import alexnet
+import torch
+import torch.nn as nn
+from torch.autograd import Variable
+import torch.nn.functional as F
+from tqdm import tqdm
+
+from data_loader import get_loader
+import config
+from utils import accuracy, Tracker
+from coral import coral
+
+
+def train(model, optimizer, source_loader, target_loader, epoch=0):
+
+    tracker = Tracker()
+    model.train()
+    tracker_class, tracker_params = tracker.MovingMeanMonitor, {'momentum': 0.99}
+
+    # Trackers to monitor classification and CORAL loss
+    classification_loss_tracker = tracker.track('classification_loss', tracker_class(**tracker_params))
+    coral_loss_tracker = tracker.track('CORAL_loss', tracker_class(**tracker_params))
+
+    min_n_batches = min(len(source_loader), len(target_loader))
+
+    tq = tqdm(range(min_n_batches), desc='{} E{:03d}'.format('Training + Adaptation', epoch), ncols=0)
+
+    for _ in tq:
+
+        source_data, source_label = next(iter(source_loader))
+        target_data, _ = next(iter(target_loader))  # Unsupervised Domain Adaptation
+
+        source_data, source_label = Variable(source_data), Variable(source_label)
+        target_data = Variable(target_data)
+
+        optimizer.zero_grad()
+
+        out_source = model(source_data)
+        out_target = model(target_data)
+
+        classification_loss = F.cross_entropy(out_source, source_label)
+
+        # This is where the magic happens
+        coral_loss = coral(out_source, out_target)
+        composite_loss = classification_loss + config.lambda_coral * coral_loss
+
+        composite_loss.backward()
+        optimizer.step()
+
+        classification_loss_tracker.append(classification_loss.item())
+        coral_loss_tracker.append(coral_loss.item())
+        fmt = '{:.4f}'.format
+        tq.set_postfix(classification_loss=fmt(classification_loss_tracker.mean.value),
+                       coral_loss=fmt(coral_loss_tracker.mean.value))
+
+
+def evaluate(model, data_loader, dataset_name, epoch=0):
+    model.eval()
+
+    tracker = Tracker()
+    tracker_class, tracker_params = tracker.MeanMonitor, {}
+    acc_tracker = tracker.track('accuracy', tracker_class(**tracker_params))
+
+    loader = tqdm(data_loader, desc='{} E{:03d}'.format('Evaluating on %s' % dataset_name, epoch), ncols=0)
+
+    accuracies = []
+    with torch.no_grad():
+        for target_data, target_label in loader:
+            target_data = Variable(target_data)
+            target_label = Variable(target_label)
+
+            output = model(target_data)
+
+            accuracies.append(accuracy(output, target_label))
+
+            acc_tracker.append(sum(accuracies)/len(accuracies))
+            fmt = '{:.4f}'.format
+            loader.set_postfix(accuracy=fmt(acc_tracker.mean.value))
+
+
+def main():
+
+    source_train_loader = get_loader(name_dataset='amazon', batch_size=config.batch_size, train=True)
+    target_train_loader = get_loader(name_dataset='webcam', batch_size=config.batch_size, train=True)
+
+    source_evaluate_loader = get_loader(name_dataset='amazon', batch_size=config.batch_size, train=False)
+    target_evaluate_loader = get_loader(name_dataset='webcam', batch_size=config.batch_size, train=False)
+
+    n_classes = len(source_train_loader.dataset.classes)
+
+    # ~ Paper : "We initialized the other layers with the parameters pre-trained on ImageNet"
+    # check https://github.com/pytorch/vision/blob/master/torchvision/models/alexnet.py
+    model = alexnet(pretrained=True)
+    # ~ Paper : The dimension of last fully connected layer (fc8) was set to the number of categories (31)
+    model.classifier[6] = nn.Linear(4096, config.n_classes)
+    # ~ Paper : and initialized with N(0, 0.005)
+    torch.nn.init.normal_(model.classifier[6].weight, mean=0, std=5e-3)
+
+    # Initialize bias to small constant number (http://cs231n.github.io/neural-networks-2/#init)
+    model.classifier[6].bias.data.fill_(0.01)
+
+    # ~ Paper : "The learning rate of fc8 is set to 10 times the other layers as it was training from scratch."
+    optimizer = torch.optim.SGD([
+        {'params':  model.features.parameters()},
+        {'params': model.classifier[:6].parameters()},
+        # fc8 -> 7th element (index 6) in the Sequential block
+        {'params': model.classifier[6].parameters(), 'lr': 10 * config.lr}
+    ], lr=config.lr, momentum=config.momentum)  # if not specified, the default lr is used
+
+    for i in range(config.epochs):
+        train(model, optimizer, source_train_loader, target_train_loader, i)
+        evaluate(model, source_evaluate_loader, 'source', i)
+        evaluate(model, target_evaluate_loader, 'target', i)
+
+
+if __name__ == '__main__':
+    main()

train.py

@@ -0,0 +1,88 @@
+import numpy as np
+import matplotlib.pyplot as plt
+import torch
+
+
+def imshow(image_tensor, mean, std, title=None):
+    """
+    Imshow for normalized Tensors.
+    Useful to visualize data from data loader
+    """
+
+    image = image_tensor.numpy().transpose((1, 2, 0))
+    image = std * image + mean
+    image = np.clip(image, 0, 1)
+    plt.imshow(image)
+    if title is not None:
+        plt.title(title)
+    plt.pause(0.001)  # pause a bit so that plots are updated
+
+
+def accuracy(output, target):
+
+    _, predicted = torch.max(output.data, 1)
+    total = target.size(0)
+    correct = (predicted == target).sum().item()
+    accuracy = correct/total
+
+    return accuracy
+
+
+class Tracker:
+
+    def __init__(self):
+        self.data = {}
+
+    def track(self, name, *monitors):
+        l = Tracker.ListStorage(monitors)
+        self.data.setdefault(name, []).append(l)
+        return l
+
+    def to_dict(self):
+        return {k: list(map(list, v)) for k, v in self.data.items()}
+
+    class ListStorage:
+        def __init__(self, monitors=[]):
+            self.data = []
+            self.monitors = monitors
+            for monitor in self.monitors:
+                setattr(self, monitor.name, monitor)
+
+        def append(self, item):
+            for monitor in self.monitors:
+                monitor.update(item)
+            self.data.append(item)
+
+        def __iter__(self):
+            return iter(self.data)
+
+    class MeanMonitor:
+        name = 'mean'
+
+        def __init__(self):
+            self.n = 0
+            self.total = 0
+
+        def update(self, value):
+            self.total += value
+            self.n += 1
+
+        @property
+        def value(self):
+            return self.total / self.n
+
+    class MovingMeanMonitor:
+        name = 'mean'
+
+        def __init__(self, momentum=0.9):
+            self.momentum = momentum
+            self.first = True
+            self.value = None
+
+        def update(self, value):
+            if self.first:
+                self.value = value
+                self.first = False
+            else:
+                m = self.momentum
+                self.value = m * self.value + (1 - m) * value