Fixes run on TPU

- applied black
- refactored trainers

Author: vfdev-5

README.md

@@ -3,6 +3,7 @@
 Experiments with "FixMatch" on Cifar10 dataset.
 
 Based on ["FixMatch: Simplifying Semi-Supervised Learning withConsistency and Confidence"](https://arxiv.org/abs/2001.07685)
+and its official [code](https://github.com/google-research/fixmatch).
 
 ## Requirements
 

ctaugment/__init__.py

@@ -5,14 +5,15 @@
 
 
 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"]])
+        return OrderedDict(
+            [(k, getattr(self, k)) for k in ["decay", "depth", "th", "rates"]]
+        )
 
 
 def get_default_cta():
@@ -32,9 +33,17 @@ def deserialize(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()))
+    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):

ctaugment/ctaugment.py

@@ -22,8 +22,8 @@
 
 
 OPS = {}
-OP = namedtuple('OP', ('f', 'bins'))
-Sample = namedtuple('Sample', ('train', 'probe'))
+OP = namedtuple("OP", ("f", "bins"))
+Sample = namedtuple("Sample", ("train", "probe"))
 
 
 def register(*bins):
@@ -41,7 +41,7 @@ def __init__(self, depth=2, th=0.85, decay=0.99):
         self.th = th
         self.rates = {}
         for k, op in OPS.items():
-            self.rates[k] = tuple([np.ones(x, 'f') for x in op.bins])
+            self.rates[k] = tuple([np.ones(x, "f") for x in op.bins])
 
     def rate_to_p(self, rate):
         p = rate + (1 - self.decay)  # Avoid to have all zero.
@@ -78,9 +78,17 @@ def update_rates(self, policy, proximity):
                 rate[p] = rate[p] * self.decay + proximity * (1 - self.decay)
 
     def stats(self):
-        return '\n'.join('%-16s    %s' % (k, ' / '.join(' '.join('%.2f' % x for x in self.rate_to_p(rate))
-                                                        for rate in self.rates[k]))
-                         for k in sorted(OPS.keys()))
+        return "\n".join(
+            "%-16s    %s"
+            % (
+                k,
+                " / ".join(
+                    " ".join("%.2f" % x for x in self.rate_to_p(rate))
+                    for rate in self.rates[k]
+                ),
+            )
+            for k in sorted(OPS.keys())
+        )
 
 
 def _enhance(x, op, level):
@@ -128,7 +136,10 @@ def cutout(x, level):
     height_loc = np.random.randint(low=0, high=img_height)
     width_loc = np.random.randint(low=0, high=img_width)
     upper_coord = (max(0, height_loc - size // 2), max(0, width_loc - size // 2))
-    lower_coord = (min(img_height, height_loc + size // 2), min(img_width, width_loc + size // 2))
+    lower_coord = (
+        min(img_height, height_loc + size // 2),
+        min(img_width, width_loc + size // 2),
+    )
     pixels = x.load()  # create the pixel map
     for i in range(upper_coord[0], lower_coord[0]):  # for every col:
         for j in range(upper_coord[1], lower_coord[1]):  # For every row
@@ -162,7 +173,14 @@ def rescale(x, scale, method):
     s = x.size
     scale *= 0.25
     crop = (scale * s[0], scale * s[1], s[0] * (1 - scale), s[1] * (1 - scale))
-    methods = (Image.ANTIALIAS, Image.BICUBIC, Image.BILINEAR, Image.BOX, Image.HAMMING, Image.NEAREST)
+    methods = (
+        Image.ANTIALIAS,
+        Image.BICUBIC,
+        Image.BILINEAR,
+        Image.BOX,
+        Image.HAMMING,
+        Image.NEAREST,
+    )
     method = methods[int(method * 5.99)]
     return x.crop(crop).resize(x.size, method)
 

dataflow/__init__.py

@@ -6,7 +6,6 @@
 
 
 class TransformedDataset(Dataset):
-
     def __init__(self, dataset, transforms):
         self.dataset = dataset
         self.transforms = transforms
@@ -31,17 +30,23 @@ def cycle(dataloader):
             yield b
 
 
-def get_supervised_train_loader(dataset_name, root, num_train_samples_per_class, download=True, **dataloader_kwargs):
+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
+        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
+            root,
+            num_train_samples_per_class=num_train_samples_per_class,
+            download=download,
         )
 
-        return get_supervised_train_loader(
-            train_dataset, **dataloader_kwargs
-        )
+        return get_supervised_train_loader(train_dataset, **dataloader_kwargs)
 
     else:
         raise ValueError("Unhandled dataset: {}".format(dataset_name))
@@ -57,7 +62,9 @@ def get_test_loader(dataset_name, root, download=True, **dataloader_kwargs):
         raise ValueError("Unhandled dataset: {}".format(dataset_name))
 
 
-def get_unsupervised_train_loader(dataset_name, root, cta, download=True, **dataloader_kwargs):
+def get_unsupervised_train_loader(
+    dataset_name, root, cta, download=True, **dataloader_kwargs
+):
     if dataset_name == "cifar10":
         from dataflow import cifar10
 
@@ -78,20 +85,24 @@ def get_unsupervised_train_loader(dataset_name, root, cta, download=True, **data
         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):
+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
+            root,
+            num_train_samples_per_class=num_train_samples_per_class,
+            download=download,
         )
 
-        return get_cta_probe_loader(
-            train_dataset,
-            cta=cta,
-            **dataloader_kwargs
-        )
+        return get_cta_probe_loader(train_dataset, cta=cta, **dataloader_kwargs)
 
     else:
         raise ValueError("Unhandled dataset: {}".format(dataset_name))
-