Initial commit

Author: liammcgold

Dataloader.py

@@ -0,0 +1,71 @@
+from torch.utils import data
+import os
+from multiprocessing import Lock
+import random
+from PIL import Image
+import numpy as np
+import torch
+
+
+class FolderDataset(data.Dataset, ):
+
+    def __init__(self, folder_loc, files=None):
+        if folder_loc[-1] != "/":
+            self.loc = folder_loc + "/"
+        else:
+            self.loc = folder_loc
+        if files is None:
+            self.files = [str(x) for x in os.listdir(folder_loc) if x.endswith('.jpg')]
+        else:
+            self.files = files
+        self.locks = [Lock() for _ in self.files]
+
+        random.Random().shuffle(self.files)
+
+    def __getitem__(self, index):
+        file = self.files[index]
+        age = int(file.split("_")[0])
+        with self.locks[index]:
+            im = Image.open(self.loc + file)
+        im = np.asarray(im)
+        return (im, age)
+
+    def __len__(self):
+        return len(self.files)
+
+        # Train Test Val
+
+
+class FolderDataloader(data.DataLoader):
+
+    def __init__(self, dataset, collate_fn=None, num_workers=2, batch_size=1):
+        super().__init__(dataset, collate_fn=collate_fn, num_workers=num_workers, batch_size=batch_size, shuffle=True)
+
+
+def collate(batch):
+    images = np.asarray([x[0] for x in batch])
+    labels = [x[1] for x in batch]
+    images = np.einsum("bxyc->bcxy", images)
+
+    return torch.tensor(images, dtype=torch.float), torch.tensor(labels, dtype=torch.uint8)
+
+
+def generate_dataloaders(folder_loc, split=(0.6, 0.3, 0.1), collate_fxn=collate, num_workers=2, batch_size=1):
+    files = [str(x) for x in os.listdir(folder_loc) if x.endswith('.jpg')]
+    random.Random(0).shuffle(files)
+    split_a = int(len(files) * split[0])
+    split_b = split_a + int(len(files) * split[1])
+    split_c = split_b + int(len(files) * split[2])
+
+    train = files[0:split_a]
+    test = files[split_a:split_b]
+    val = files[split_b:]
+
+    train_ds = FolderDataset(folder_loc, files=train)
+    train_dl = FolderDataloader(train_ds, collate_fn=collate_fxn, num_workers=num_workers, batch_size=batch_size)
+    test_ds = FolderDataset(folder_loc, files=test)
+    test_dl = FolderDataloader(test_ds, collate_fn=collate_fxn, num_workers=num_workers)
+    val_ds = FolderDataset(folder_loc, files=val)
+    val_dl = FolderDataloader(val_ds, collate_fn=collate_fxn, num_workers=num_workers)
+
+    return train_dl, test_dl, val_dl

Logging.py

@@ -0,0 +1,15 @@
+def log(folder, split, batch_norm, probabilistic, epochs, output_folder, lr, batch_size):
+    params = [str(lr),
+              str(probabilistic),
+              str(split),
+              str(batch_norm),
+              str(epochs),
+              str(batch_size),
+              str(folder),
+              str(output_folder)]
+    log = ""
+    for entry in params:
+        log += "," + entry
+
+    with open(output_folder + "Description.txt", "w") as file:
+        file.write(log)

Losses.py

@@ -0,0 +1,12 @@
+from torch import distributions as d
+from torch.nn import functional as F
+
+
+def reg_loss(x, y):
+    target = y.float().unsqueeze(-1).to('cuda')
+    return F.mse_loss(x, target)
+
+
+def prob_loss(x, y):
+    n = d.Normal(x[0][0], x[0][1])
+    return -1 * n.log_prob(y)

Models.py

@@ -0,0 +1,88 @@
+import torch as t
+from torch import nn
+
+
+class Regress(nn.Module):
+    def __init__(self, probabilistic=False, in_chan=3, batch_norm=False):
+        super(Regress, self).__init__()
+        n_out = 1 + int(probabilistic)
+        self.conv_blocks = nn.Sequential(
+            ResMod(in_chan=in_chan, out_chan=in_chan, batch_norm=batch_norm),
+            ResMod(in_chan=in_chan, out_chan=in_chan, batch_norm=batch_norm),
+            nn.Conv2d(in_chan, 16, 3, stride=2),
+            nn.Conv2d(16, 32, 3, stride=2),
+            nn.Conv2d(32, 64, 3, stride=2),
+            nn.Conv2d(64, 128, 3, stride=2),
+        )
+        if probabilistic:
+            s = 210
+        else:
+            s = 190
+        adjustment = s // 2 ** 4
+
+        # Output layers
+        self.dense_unit = nn.Sequential(nn.Linear(128 * adjustment ** 2, n_out * 16),
+                                        nn.Linear(n_out * 16, n_out * 8),
+                                        nn.Linear(n_out * 8, n_out * 4),
+                                        nn.Linear(n_out * 4, n_out * 2),
+                                        nn.Linear(n_out * 2, n_out))
+
+        self.frozen = False
+
+    def forward(self, img):
+        out = self.conv_blocks(img)
+        out = out.view(out.shape[0], -1)
+        y = self.dense_unit(out)
+        return y
+
+
+class ResMod(nn.Module):
+    def __init__(self, batch_norm=False, in_chan=1, out_chan=1):
+        super(ResMod, self).__init__()
+
+        self.in_chan = in_chan
+        self.out_chan = out_chan
+
+        ###############
+        # first block #
+        ###############
+        module = nn.ModuleList()
+
+        module.append(nn.Conv2d(in_chan, out_chan, (3, 3), padding=1))
+
+        if batch_norm:
+            module.append(nn.BatchNorm2d(out_chan))
+
+        module.append(nn.ReLU())
+
+        self.first_block = module
+
+        ################
+        # second block #
+        ################
+        module = nn.ModuleList()
+
+        module.append(nn.Conv2d(out_chan, out_chan, (3, 3), padding=1))
+
+        if batch_norm:
+            module.append(nn.BatchNorm2d(out_chan))
+
+        self.second_block = module
+
+        self.shortcut = nn.Conv2d(in_chan, out_chan, (1, 1))
+
+        self.relu = nn.ReLU()
+
+    def forward(self, x):
+        residual = x
+
+        for module in self.first_block:
+            x = module(x)
+        for module in self.second_block:
+            x = module(x)
+
+        x += self.shortcut(residual)
+
+        x = self.relu(x)
+
+        return x