[WIP] fully supervised training app

Author: vfdev-5

README.md

@@ -7,6 +7,7 @@ Based on ["FixMatch: Simplifying Semi-Supervised Learning withConsistency and Co
 ## Requirements
 
 ```bash
+pip install --upgrade --pre hydra-core
 pip install --upgrade --pre pytorch-ignite
 ```
 

config/dataflow/cifar10.yaml

@@ -0,0 +1,7 @@
+# @package _group_
+name: cifar10
+
+data_path: "/tmp/cifar10"
+
+batch_size: 64
+num_workers: 12

config/fixmatch.yaml

@@ -0,0 +1,15 @@
+hydra:
+  run:
+    dir: /tmp/output-fixmatch-cifar10-hydra/fully_supervised/${now:%Y%m%d-%H%M%S}
+
+seed: 543
+model: "resnet18"
+
+
+
+defaults:
+  - dataset: cifar10
+  - solver: default
+
+
+

config/fully_supervised.yaml

@@ -0,0 +1,36 @@
+hydra:
+  run:
+    dir: /tmp/output-fixmatch-cifar10-hydra/fully_supervised/${now:%Y%m%d-%H%M%S}
+  job_logging:
+    handlers:
+      console:
+        level: WARN
+    root:
+      level: WARN
+
+name: fully-supervised
+
+seed: 543
+debug: false
+
+# model name (from torchvision) to setup model to train. For Wide-Resnet, use "WRN-28-2"
+model: "resnet18"
+num_classes: 10
+
+ema_decay: 0.999
+
+defaults:
+  - dataflow: cifar10
+  - solver: default
+  - ssl: full_sup
+
+
+distributed:
+  # backend to use for distributed configuration. Possible values: None, "nccl", "xla-tpu", "gloo" etc. Default, None.
+  backend: null
+  # optional argument to setup number of processes per node. It is useful, when main python process is spawning training as child processes.
+  nproc_per_node: null
+
+
+online_exp_tracking:
+  wandb: false

config/solver/default.yaml

@@ -0,0 +1,30 @@
+# @package _group_
+
+num_epochs: 1024
+
+epoch_length: 128  # epoch_length * num_epochs == 2 ** 20
+
+checkpoint_every: 500
+
+validate_every: 1
+
+resume_from: null
+
+optimizer:
+  cls: torch.optim.SGD
+  params:
+    lr: 0.01
+    momentum: 0.9
+    weight_decay: 0.0001
+    nesterov: false
+
+
+supervised_criterion:
+  cls: torch.nn.CrossEntropyLoss
+
+
+lr_scheduler:
+  cls: torch.optim.lr_scheduler.CosineAnnealingLR
+  params:
+    eta_min: 0.0
+    T_max: null

config/ssl/full_sup.yaml

@@ -0,0 +1,3 @@
+# @package _group_
+
+num_train_samples_per_class: 25

configs.py

@@ -28,13 +28,15 @@ def get_default_config():
         "model": "WRN-28-2",
         "momentum": 0.9,
         "weight_decay": 0.0005,
+
         "batch_size": batch_size,
         "num_workers": 12,
+
         "num_epochs": 1024,
         "epoch_length": 2 ** 16 // batch_size,  # epoch_length * num_epochs == 2 ** 20
         "learning_rate": 0.03,
-        "validate_every": 1,
 
+        "validate_every": 1,
         # Logging:
         "display_iters": True,
         "checkpoint_every": 200,

ctaugment/__init__.py

@@ -0,0 +1,53 @@
+import json
+from collections import OrderedDict
+
+from ctaugment.ctaugment import *
+
+
+class StorableCTAugment(CTAugment):
+
+    def load_state_dict(self, state):
+        for k in ["decay", "depth", "th", "rates"]:
+            assert k in state, "{} not in {}".format(k, state.keys())
+            setattr(self, k, state[k])
+
+    def state_dict(self):
+        return OrderedDict([(k, getattr(self, k)) for k in ["decay", "depth", "th", "rates"]])
+
+
+def get_default_cta():
+    return StorableCTAugment()
+
+
+def cta_apply(pil_img, ops):
+    if ops is None:
+        return pil_img
+    for op, args in ops:
+        pil_img = OPS[op].f(pil_img, *args)
+    return pil_img
+
+
+def deserialize(policy_str):
+    return [OP(f=x[0], bins=x[1]) for x in json.loads(policy_str)]
+
+
+def stats(cta):
+    return '\n'.join('%-16s    %s' % (k, ' / '.join(' '.join('%.2f' % x for x in cta.rate_to_p(rate))
+                                                    for rate in cta.rates[k]))
+                     for k in sorted(OPS.keys()))
+
+
+def interleave(x, batch, inverse=False):
+    """
+    TF code
+    def interleave(x, batch):
+        s = x.get_shape().as_list()
+        return tf.reshape(tf.transpose(tf.reshape(x, [-1, batch] + s[1:]), [1, 0] + list(range(2, 1+len(s)))), [-1] + s[1:])
+    """
+    shape = x.shape
+    axes = [batch, -1] if inverse else [-1, batch]
+    return x.reshape(*axes, *shape[1:]).transpose(0, 1).reshape(-1, *shape[1:])
+
+
+def deinterleave(x, batch):
+    return interleave(x, batch, inverse=True)

ctaugment.py renamed to ctaugment/ctaugment.py

@@ -0,0 +1,97 @@
+from functools import partial
+
+from torch.utils.data import Dataset
+
+from ignite.utils import convert_tensor
+
+
+class TransformedDataset(Dataset):
+
+    def __init__(self, dataset, transforms):
+        self.dataset = dataset
+        self.transforms = transforms
+
+    def __getitem__(self, i):
+        dp = self.dataset[i]
+        return self.transforms(dp)
+
+    def __len__(self):
+        return len(self.dataset)
+
+
+def sup_prepare_batch(batch, device, non_blocking):
+    x = convert_tensor(batch["image"], device, non_blocking)
+    y = convert_tensor(batch["target"], device, non_blocking)
+    return x, y
+
+
+def cycle(dataloader):
+    while True:
+        for b in dataloader:
+            yield b
+
+
+def get_supervised_train_loader(dataset_name, root, num_train_samples_per_class, download=True, **dataloader_kwargs):
+    if dataset_name == "cifar10":
+        from dataflow.cifar10 import get_supervised_trainset, get_supervised_train_loader, weak_transforms
+
+        train_dataset = get_supervised_trainset(
+            root, num_train_samples_per_class=num_train_samples_per_class, download=download
+        )
+
+        return get_supervised_train_loader(
+            train_dataset, **dataloader_kwargs
+        )
+
+    else:
+        raise ValueError("Unhandled dataset: {}".format(dataset_name))
+
+
+def get_test_loader(dataset_name, root, download=True, **dataloader_kwargs):
+    if dataset_name == "cifar10":
+        from dataflow.cifar10 import get_test_loader
+
+        return get_test_loader(root=root, download=download, **dataloader_kwargs)
+
+    else:
+        raise ValueError("Unhandled dataset: {}".format(dataset_name))
+
+
+def get_unsupervised_train_loader(dataset_name, root, cta, download=True, **dataloader_kwargs):
+    if dataset_name == "cifar10":
+        from dataflow import cifar10
+
+        full_train_dataset = cifar10.get_supervised_trainset(
+            root, num_train_samples_per_class=None, download=download
+        )
+
+        strong_transforms = partial(cifar10.cta_image_transforms, cta=cta)
+
+        return get_unsupervised_train_loader(
+            full_train_dataset,
+            transforms_weak=cifar10.weak_transforms,
+            transforms_strong=strong_transforms,
+            **dataloader_kwargs
+        )
+
+    else:
+        raise ValueError("Unhandled dataset: {}".format(dataset_name))
+
+
+def get_cta_probe_loader(dataset_name, root, num_train_samples_per_class, cta, download=True, **dataloader_kwargs):
+    if dataset_name == "cifar10":
+        from dataflow.cifar10 import get_supervised_trainset, get_cta_probe_loader
+
+        train_dataset = get_supervised_trainset(
+            root, num_train_samples_per_class=num_train_samples_per_class, download=download
+        )
+
+        return get_cta_probe_loader(
+            train_dataset,
+            cta=cta,
+            **dataloader_kwargs
+        )
+
+    else:
+        raise ValueError("Unhandled dataset: {}".format(dataset_name))
+