From 780d90e8dd1779c2746378210688afa602df2adf Mon Sep 17 00:00:00 2001
From: Nathan Michlo <NathanJMichlo@gmail.com>
Date: Sat, 6 Aug 2022 00:22:09 +0200
Subject: [PATCH 1/4] fix examples with `num_workers` in DataLoaders

---
 README.md                       | 39 +++++++++++++++++----------------
 docs/examples/mnist_example.py  |  6 +++--
 docs/examples/readme_example.py | 39 +++++++++++++++++----------------
 3 files changed, 44 insertions(+), 40 deletions(-)

diff --git a/README.md b/README.md
index 241bd4f9..a086897f 100644
--- a/README.md
+++ b/README.md
@@ -315,7 +315,6 @@ beta schedule and evaluates the trained model with various metrics.
 <p>
 
 ```python3
-import os
 import pytorch_lightning as pl
 import torch
 from torch.utils.data import DataLoader
@@ -334,39 +333,41 @@ from disent.schedule import CyclicSchedule
 
 # create the dataset & dataloaders
 # - ToImgTensorF32 transforms images from numpy arrays to tensors and performs checks
+# - if you use `num_workers` in the DataLoader, the make sure to wrap `trainer.fit`
+#   with `if __name__ == '__main__': ...`
 data = XYObjectData()
 dataset = DisentDataset(dataset=data, sampler=SingleSampler(), transform=ToImgTensorF32())
-dataloader = DataLoader(dataset=dataset, batch_size=128, shuffle=True, num_workers=os.cpu_count())
+dataloader = DataLoader(dataset=dataset, batch_size=128, shuffle=True)
 
 # create the BetaVAE model
 # - adjusting the beta, learning rate, and representation size.
 module = BetaVae(
-  model=AutoEncoder(
-    # z_multiplier is needed to output mu & logvar when parameterising normal distribution
-    encoder=EncoderConv64(x_shape=data.x_shape, z_size=10, z_multiplier=2),
-    decoder=DecoderConv64(x_shape=data.x_shape, z_size=10),
-  ),
-  cfg=BetaVae.cfg(
-    optimizer='adam',
-    optimizer_kwargs=dict(lr=1e-3),
-    loss_reduction='mean_sum',
-    beta=4,
-  )
+    model=AutoEncoder(
+        # z_multiplier is needed to output mu & logvar when parameterising normal distribution
+        encoder=EncoderConv64(x_shape=data.x_shape, z_size=10, z_multiplier=2),
+        decoder=DecoderConv64(x_shape=data.x_shape, z_size=10),
+    ),
+    cfg=BetaVae.cfg(
+        optimizer='adam',
+        optimizer_kwargs=dict(lr=1e-3),
+        loss_reduction='mean_sum',
+        beta=4,
+    )
 )
 
 # cyclic schedule for target 'beta' in the config/cfg. The initial value from the
 # config is saved and multiplied by the ratio from the schedule on each step.
 # - based on: https://arxiv.org/abs/1903.10145
 module.register_schedule(
-  'beta', CyclicSchedule(
-    period=1024,  # repeat every: trainer.global_step % period
-  )
+    'beta', CyclicSchedule(
+        period=1024,  # repeat every: trainer.global_step % period
+    )
 )
 
 # train model
 # - for 2048 batches/steps
 trainer = pl.Trainer(
-  max_steps=2048, gpus=1 if torch.cuda.is_available() else None, logger=False, checkpoint_callback=False
+    max_steps=2048, gpus=1 if torch.cuda.is_available() else None, logger=False, checkpoint_callback=False
 )
 trainer.fit(module, dataloader)
 
@@ -376,8 +377,8 @@ trainer.fit(module, dataloader)
 get_repr = lambda x: module.encode(x.to(module.device))
 
 metrics = {
-  **metric_dci(dataset, get_repr, num_train=1000, num_test=500, show_progress=True),
-  **metric_mig(dataset, get_repr, num_train=2000),
+    **metric_dci(dataset, get_repr, num_train=1000, num_test=500, show_progress=True),
+    **metric_mig(dataset, get_repr, num_train=2000),
 }
 
 # evaluate
diff --git a/docs/examples/mnist_example.py b/docs/examples/mnist_example.py
index 6e56354d..caae3bc9 100644
--- a/docs/examples/mnist_example.py
+++ b/docs/examples/mnist_example.py
@@ -24,8 +24,10 @@ def __getitem__(self, index):
 dataset_test  =               MNIST(data_folder, train=False, download=True, transform=ToImgTensorF32())
 
 # create the dataloaders
-dataloader_train = DataLoader(dataset=dataset_train, batch_size=128, shuffle=True, num_workers=os.cpu_count())
-dataloader_test  = DataLoader(dataset=dataset_test,  batch_size=128, shuffle=True, num_workers=os.cpu_count())
+# - if you use `num_workers` in the DataLoader, the make sure to wrap `trainer.fit`
+#   with `if __name__ == '__main__': ...`
+dataloader_train = DataLoader(dataset=dataset_train, batch_size=128, shuffle=True)
+dataloader_test  = DataLoader(dataset=dataset_test,  batch_size=128, shuffle=True)
 
 # create the model
 module = AdaVae(
diff --git a/docs/examples/readme_example.py b/docs/examples/readme_example.py
index 40d04007..ab794444 100644
--- a/docs/examples/readme_example.py
+++ b/docs/examples/readme_example.py
@@ -1,4 +1,3 @@
-import os
 import pytorch_lightning as pl
 import torch
 from torch.utils.data import DataLoader
@@ -17,39 +16,41 @@
 
 # create the dataset & dataloaders
 # - ToImgTensorF32 transforms images from numpy arrays to tensors and performs checks
+# - if you use `num_workers` in the DataLoader, the make sure to wrap `trainer.fit`
+#   with `if __name__ == '__main__': ...`
 data = XYObjectData()
 dataset = DisentDataset(dataset=data, sampler=SingleSampler(), transform=ToImgTensorF32())
-dataloader = DataLoader(dataset=dataset, batch_size=128, shuffle=True, num_workers=os.cpu_count())
+dataloader = DataLoader(dataset=dataset, batch_size=128, shuffle=True)
 
 # create the BetaVAE model
 # - adjusting the beta, learning rate, and representation size.
 module = BetaVae(
-  model=AutoEncoder(
-    # z_multiplier is needed to output mu & logvar when parameterising normal distribution
-    encoder=EncoderConv64(x_shape=data.x_shape, z_size=10, z_multiplier=2),
-    decoder=DecoderConv64(x_shape=data.x_shape, z_size=10),
-  ),
-  cfg=BetaVae.cfg(
-    optimizer='adam',
-    optimizer_kwargs=dict(lr=1e-3),
-    loss_reduction='mean_sum',
-    beta=4,
-  )
+    model=AutoEncoder(
+        # z_multiplier is needed to output mu & logvar when parameterising normal distribution
+        encoder=EncoderConv64(x_shape=data.x_shape, z_size=10, z_multiplier=2),
+        decoder=DecoderConv64(x_shape=data.x_shape, z_size=10),
+    ),
+    cfg=BetaVae.cfg(
+        optimizer='adam',
+        optimizer_kwargs=dict(lr=1e-3),
+        loss_reduction='mean_sum',
+        beta=4,
+    )
 )
 
 # cyclic schedule for target 'beta' in the config/cfg. The initial value from the
 # config is saved and multiplied by the ratio from the schedule on each step.
 # - based on: https://arxiv.org/abs/1903.10145
 module.register_schedule(
-  'beta', CyclicSchedule(
-    period=1024,  # repeat every: trainer.global_step % period
-  )
+    'beta', CyclicSchedule(
+        period=1024,  # repeat every: trainer.global_step % period
+    )
 )
 
 # train model
 # - for 2048 batches/steps
 trainer = pl.Trainer(
-  max_steps=2048, gpus=1 if torch.cuda.is_available() else None, logger=False, checkpoint_callback=False
+    max_steps=2048, gpus=1 if torch.cuda.is_available() else None, logger=False, checkpoint_callback=False
 )
 trainer.fit(module, dataloader)
 
@@ -59,8 +60,8 @@
 get_repr = lambda x: module.encode(x.to(module.device))
 
 metrics = {
-  **metric_dci(dataset, get_repr, num_train=1000, num_test=500, show_progress=True),
-  **metric_mig(dataset, get_repr, num_train=2000),
+    **metric_dci(dataset, get_repr, num_train=1000, num_test=500, show_progress=True),
+    **metric_mig(dataset, get_repr, num_train=2000),
 }
 
 # evaluate

From a51972d7fc16be224ff5b73db724a84c692759aa Mon Sep 17 00:00:00 2001
From: Nathan Michlo <NathanJMichlo@gmail.com>
Date: Sat, 6 Aug 2022 00:23:51 +0200
Subject: [PATCH 2/4] fix requirements

---
 requirements.txt | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/requirements.txt b/requirements.txt
index 6cc726f0..8d527fb7 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -6,7 +6,7 @@ numpy>=1.19.0
 torch>=1.9.0
 torchvision>=0.10.0
 pytorch-lightning>=1.4.0,<1.7
-torch_optimizer>=0.1.0
+torch_optimizer>=0.1.0,!=0.2
 scipy>=1.7.0
 scikit-learn>=0.24.2
 

From e70867cada0b5a62d92d0a6a316866c6bc5595ca Mon Sep 17 00:00:00 2001
From: Nathan Michlo <NathanJMichlo@gmail.com>
Date: Sat, 6 Aug 2022 00:36:33 +0200
Subject: [PATCH 3/4] explicitly set num_workers=0 in examples

---
 README.md                                             | 6 +++---
 docs/examples/mnist_example.py                        | 8 ++++----
 docs/examples/overview_dataset_loader.py              | 2 +-
 docs/examples/overview_framework_adagvae.py           | 2 +-
 docs/examples/overview_framework_ae.py                | 2 +-
 docs/examples/overview_framework_betavae.py           | 2 +-
 docs/examples/overview_framework_betavae_scheduled.py | 2 +-
 docs/examples/overview_metrics.py                     | 2 +-
 docs/examples/readme_example.py                       | 6 +++---
 tests/test_frameworks.py                              | 4 ++--
 10 files changed, 18 insertions(+), 18 deletions(-)

diff --git a/README.md b/README.md
index a086897f..b7cbdbaf 100644
--- a/README.md
+++ b/README.md
@@ -333,11 +333,11 @@ from disent.schedule import CyclicSchedule
 
 # create the dataset & dataloaders
 # - ToImgTensorF32 transforms images from numpy arrays to tensors and performs checks
-# - if you use `num_workers` in the DataLoader, the make sure to wrap `trainer.fit`
-#   with `if __name__ == '__main__': ...`
+# - if you use `num_workers != 0` in the DataLoader, the make sure to
+#   wrap `trainer.fit` with `if __name__ == '__main__': ...`
 data = XYObjectData()
 dataset = DisentDataset(dataset=data, sampler=SingleSampler(), transform=ToImgTensorF32())
-dataloader = DataLoader(dataset=dataset, batch_size=128, shuffle=True)
+dataloader = DataLoader(dataset=dataset, batch_size=128, shuffle=True, num_workers=0)
 
 # create the BetaVAE model
 # - adjusting the beta, learning rate, and representation size.
diff --git a/docs/examples/mnist_example.py b/docs/examples/mnist_example.py
index caae3bc9..3151c13f 100644
--- a/docs/examples/mnist_example.py
+++ b/docs/examples/mnist_example.py
@@ -24,10 +24,10 @@ def __getitem__(self, index):
 dataset_test  =               MNIST(data_folder, train=False, download=True, transform=ToImgTensorF32())
 
 # create the dataloaders
-# - if you use `num_workers` in the DataLoader, the make sure to wrap `trainer.fit`
-#   with `if __name__ == '__main__': ...`
-dataloader_train = DataLoader(dataset=dataset_train, batch_size=128, shuffle=True)
-dataloader_test  = DataLoader(dataset=dataset_test,  batch_size=128, shuffle=True)
+# - if you use `num_workers != 0` in the DataLoader, the make sure to
+#   wrap `trainer.fit` with `if __name__ == '__main__': ...`
+dataloader_train = DataLoader(dataset=dataset_train, batch_size=128, shuffle=True, num_workers=0)
+dataloader_test  = DataLoader(dataset=dataset_test,  batch_size=128, shuffle=True, num_workers=0)
 
 # create the model
 module = AdaVae(
diff --git a/docs/examples/overview_dataset_loader.py b/docs/examples/overview_dataset_loader.py
index 21aca53c..75d35e85 100644
--- a/docs/examples/overview_dataset_loader.py
+++ b/docs/examples/overview_dataset_loader.py
@@ -7,7 +7,7 @@
 # prepare the data
 data = XYObjectData(grid_size=4, min_square_size=1, max_square_size=2, square_size_spacing=1, palette='rgb_1')
 dataset = DisentDataset(data, sampler=GroundTruthPairOrigSampler(), transform=ToImgTensorF32())
-dataloader = DataLoader(dataset=dataset, batch_size=4, shuffle=True)
+dataloader = DataLoader(dataset=dataset, batch_size=4, shuffle=True, num_workers=0)
 
 # iterate over single epoch
 for batch in dataloader:
diff --git a/docs/examples/overview_framework_adagvae.py b/docs/examples/overview_framework_adagvae.py
index e0d0d530..883a7e9f 100644
--- a/docs/examples/overview_framework_adagvae.py
+++ b/docs/examples/overview_framework_adagvae.py
@@ -13,7 +13,7 @@
 # prepare the data
 data = XYObjectData()
 dataset = DisentDataset(data, GroundTruthPairOrigSampler(), transform=ToImgTensorF32())
-dataloader = DataLoader(dataset=dataset, batch_size=4, shuffle=True)
+dataloader = DataLoader(dataset=dataset, batch_size=4, shuffle=True, num_workers=0)
 
 # create the pytorch lightning system
 module: pl.LightningModule = AdaVae(
diff --git a/docs/examples/overview_framework_ae.py b/docs/examples/overview_framework_ae.py
index f848e305..1a108c2a 100644
--- a/docs/examples/overview_framework_ae.py
+++ b/docs/examples/overview_framework_ae.py
@@ -12,7 +12,7 @@
 # prepare the data
 data = XYObjectData()
 dataset = DisentDataset(data, transform=ToImgTensorF32())
-dataloader = DataLoader(dataset=dataset, batch_size=4, shuffle=True)
+dataloader = DataLoader(dataset=dataset, batch_size=4, shuffle=True, num_workers=0)
 
 # create the pytorch lightning system
 module: pl.LightningModule = Ae(
diff --git a/docs/examples/overview_framework_betavae.py b/docs/examples/overview_framework_betavae.py
index 478eb421..5019c6fd 100644
--- a/docs/examples/overview_framework_betavae.py
+++ b/docs/examples/overview_framework_betavae.py
@@ -12,7 +12,7 @@
 # prepare the data
 data = XYObjectData()
 dataset = DisentDataset(data, transform=ToImgTensorF32())
-dataloader = DataLoader(dataset=dataset, batch_size=4, shuffle=True)
+dataloader = DataLoader(dataset=dataset, batch_size=4, shuffle=True, num_workers=0)
 
 # create the pytorch lightning system
 module: pl.LightningModule = BetaVae(
diff --git a/docs/examples/overview_framework_betavae_scheduled.py b/docs/examples/overview_framework_betavae_scheduled.py
index 290d55b2..4a26cc37 100644
--- a/docs/examples/overview_framework_betavae_scheduled.py
+++ b/docs/examples/overview_framework_betavae_scheduled.py
@@ -12,7 +12,7 @@
 # prepare the data
 data = XYObjectData()
 dataset = DisentDataset(data, transform=ToImgTensorF32())
-dataloader = DataLoader(dataset=dataset, batch_size=4, shuffle=True)
+dataloader = DataLoader(dataset=dataset, batch_size=4, shuffle=True, num_workers=0)
 
 # create the pytorch lightning system
 module: pl.LightningModule = BetaVae(
diff --git a/docs/examples/overview_metrics.py b/docs/examples/overview_metrics.py
index cd83103f..871971cc 100644
--- a/docs/examples/overview_metrics.py
+++ b/docs/examples/overview_metrics.py
@@ -11,7 +11,7 @@
 
 data = XYObjectData()
 dataset = DisentDataset(data, transform=ToImgTensorF32(), augment=None)
-dataloader = DataLoader(dataset=dataset, batch_size=32, shuffle=True)
+dataloader = DataLoader(dataset=dataset, batch_size=32, shuffle=True, num_workers=0)
 
 def make_vae(beta):
     return BetaVae(
diff --git a/docs/examples/readme_example.py b/docs/examples/readme_example.py
index ab794444..60b5a0b8 100644
--- a/docs/examples/readme_example.py
+++ b/docs/examples/readme_example.py
@@ -16,11 +16,11 @@
 
 # create the dataset & dataloaders
 # - ToImgTensorF32 transforms images from numpy arrays to tensors and performs checks
-# - if you use `num_workers` in the DataLoader, the make sure to wrap `trainer.fit`
-#   with `if __name__ == '__main__': ...`
+# - if you use `num_workers != 0` in the DataLoader, the make sure to
+#   wrap `trainer.fit` with `if __name__ == '__main__': ...`
 data = XYObjectData()
 dataset = DisentDataset(dataset=data, sampler=SingleSampler(), transform=ToImgTensorF32())
-dataloader = DataLoader(dataset=dataset, batch_size=128, shuffle=True)
+dataloader = DataLoader(dataset=dataset, batch_size=128, shuffle=True, num_workers=0)
 
 # create the BetaVAE model
 # - adjusting the beta, learning rate, and representation size.
diff --git a/tests/test_frameworks.py b/tests/test_frameworks.py
index 7e3c62d3..21560f8a 100644
--- a/tests/test_frameworks.py
+++ b/tests/test_frameworks.py
@@ -105,7 +105,7 @@ def test_frameworks(Framework, cfg_kwargs, Data):
 
     data = XYObjectData() if (Data is None) else Data()
     dataset = DisentDataset(data, DataSampler(), transform=ToImgTensorF32())
-    dataloader = DataLoader(dataset=dataset, batch_size=4, shuffle=True)
+    dataloader = DataLoader(dataset=dataset, batch_size=4, shuffle=True, num_workers=0)
 
     framework = Framework(
         model=AutoEncoder(
@@ -176,7 +176,7 @@ def test_ada_vae_similarity():
 
     data = XYObjectData()
     dataset = DisentDataset(data, sampler=RandomSampler(num_samples=2), transform=ToImgTensorF32())
-    dataloader = DataLoader(dataset, num_workers=0, batch_size=3)
+    dataloader = DataLoader(dataset, batch_size=3, num_workers=0)
 
     model = AutoEncoder(
         encoder=EncoderLinear(x_shape=data.x_shape, z_size=25, z_multiplier=2),

From 77106e92a3954edc9d8e3e2a52db7fa848e8ee0c Mon Sep 17 00:00:00 2001
From: Nathan Michlo <NathanJMichlo@gmail.com>
Date: Sat, 6 Aug 2022 00:24:17 +0200
Subject: [PATCH 4/4] version bump v0.6.2

---
 setup.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/setup.py b/setup.py
index 1c6755d0..6649a77c 100644
--- a/setup.py
+++ b/setup.py
@@ -48,7 +48,7 @@
     author="Nathan Juraj Michlo",
     author_email="NathanJMichlo@gmail.com",
 
-    version="0.6.1",
+    version="0.6.2",
     python_requires=">=3.8",  # we make use of standard library features only in 3.8
     packages=setuptools.find_packages(),