Merge branch 'dev' into main

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 != 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, num_workers=os.cpu_count())
+dataloader = DataLoader(dataset=dataset, batch_size=128, shuffle=True, num_workers=0)
 
 # 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

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 != 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(

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:

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(

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(

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(

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(

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(

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 != 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, num_workers=os.cpu_count())
+dataloader = DataLoader(dataset=dataset, batch_size=128, shuffle=True, num_workers=0)
 
 # 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

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
 

setup.py

@@ -48,7 +48,7 @@
     author="Nathan Juraj Michlo",
     author_email="[email protected]",
 
-    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(),
 

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),