House-keeping

src/training/__init__.py

@@ -1,3 +1,4 @@
 from .histogram_dataset import HistogramDataset
-from .random_edit import random_edit
-from .progressive_pooling_loss import ProgressivePoolingLoss
+from .get_next_run_name import get_next_run_name
+from .random_hparam_search import random_hparam_search
+from .train import train

src/training/progressive_pooling_loss.py

@@ -1,38 +0,0 @@
-from typing import List
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-
-class ProgressivePoolingLoss(nn.Module):
-    def __init__(self, target_sizes: List[int], damping: float):
-        super(ProgressivePoolingLoss, self).__init__()
-        self._target_sizes = target_sizes
-        self._damping = damping
-
-    def forward(self, tensor_a, tensor_b):
-        assert (
-            tensor_a.size() == tensor_b.size()
-        ), f"Input tensors must have the same size, got {tensor_a.size()} and {tensor_b.size()}"
-
-        assert (
-            len(tensor_a.size()) == 5
-        ), f"Input tensors must have 5 dimensions, got {tensor_a.size()}"
-
-        _minibatch_size, _channels, depth, height, width = tensor_a.size()
-        assert depth == height == width, "Input tensors must be cubes."
-
-        loss = 0.0
-        weight = 1
-
-        for target_size in self._target_sizes:
-            pool_size = depth // target_size
-            pooled_a = F.avg_pool3d(tensor_a, pool_size) * (pool_size**3)
-            pooled_b = F.avg_pool3d(tensor_b, pool_size) * (pool_size**3)
-
-            diff = torch.abs(pooled_a - pooled_b)
-
-            loss += diff.mean() * weight
-            weight *= self._damping
-
-        return loss

src/training/random_edit.py

@@ -1,33 +0,0 @@
-from PIL import Image, ImageEnhance
-from utils import random
-from operations import (
-    add_noise,
-    add_random_colour_spill,
-    get_random_gamma,
-    adjust_gamma,
-    apply_pixel_shader,
-    get_random_brightness_lut,
-    get_random_saturation_per_hue_lut,
-)
-import numpy as np
-
-
-def random_edit(img: Image, seed: int = 42) -> Image:
-    np.random.seed(seed)
-    img = img.convert("RGB")
-
-    img = ImageEnhance.Contrast(img).enhance(random(0.5, 1.5))
-    img = adjust_gamma(img, get_random_gamma())
-
-    img = img.convert("HSV")
-    saturation_lut = get_random_saturation_per_hue_lut()
-    brightness_lut = get_random_brightness_lut()
-    img = apply_pixel_shader(
-        img, lambda h, s, v: (h, round(s * saturation_lut[h]), brightness_lut[v])
-    )
-    img = img.convert("RGB")
-
-    img = add_random_colour_spill(img, 0.2)
-    img = add_noise(img, random(0, 0.1))
-
-    return img