download

Author: Borda

pl_examples/__init__.py

@@ -1,4 +1,5 @@
 import os
+from urllib.error import HTTPError
 
 from six.moves import urllib
 
@@ -14,8 +15,16 @@
 _DATASETS_PATH = os.path.join(_PACKAGE_ROOT, 'Datasets')
 
 _TORCHVISION_AVAILABLE = _module_available("torchvision")
+_TORCHVISION_MNIST_AVAILABLE = True
 _DALI_AVAILABLE = _module_available("nvidia.dali")
 
+if _TORCHVISION_AVAILABLE:
+    try:
+        from torchvision.datasets.mnist import MNIST
+        MNIST(_DATASETS_PATH, download=True)
+    except HTTPError:
+        _TORCHVISION_MNIST_AVAILABLE = False
+
 LIGHTNING_LOGO = """
                     ####
                 ###########

pl_examples/basic_examples/autoencoder.py

@@ -20,9 +20,9 @@
 from torch.utils.data import DataLoader, random_split
 
 import pytorch_lightning as pl
-from pl_examples import _DATASETS_PATH, _TORCHVISION_AVAILABLE, cli_lightning_logo
+from pl_examples import _DATASETS_PATH, _TORCHVISION_AVAILABLE, _TORCHVISION_MNIST_AVAILABLE, cli_lightning_logo
 
-if _TORCHVISION_AVAILABLE:
+if _TORCHVISION_AVAILABLE and _TORCHVISION_MNIST_AVAILABLE:
     from torchvision import transforms
     from torchvision.datasets.mnist import MNIST
 else:

pl_examples/basic_examples/backbone_image_classifier.py

@@ -19,9 +19,9 @@
 from torch.utils.data import DataLoader, random_split
 
 import pytorch_lightning as pl
-from pl_examples import _DATASETS_PATH, _TORCHVISION_AVAILABLE, cli_lightning_logo
+from pl_examples import _DATASETS_PATH, _TORCHVISION_AVAILABLE, _TORCHVISION_MNIST_AVAILABLE, cli_lightning_logo
 
-if _TORCHVISION_AVAILABLE:
+if _TORCHVISION_AVAILABLE and _TORCHVISION_MNIST_AVAILABLE:
     from torchvision import transforms
     from torchvision.datasets.mnist import MNIST
 else:

pl_examples/basic_examples/dali_image_classifier.py

@@ -23,9 +23,15 @@
 from torch.utils.data import random_split
 
 import pytorch_lightning as pl
-from pl_examples import _DALI_AVAILABLE, _DATASETS_PATH, _TORCHVISION_AVAILABLE, cli_lightning_logo
-
-if _TORCHVISION_AVAILABLE:
+from pl_examples import (
+    _DALI_AVAILABLE,
+    _DATASETS_PATH,
+    _TORCHVISION_AVAILABLE,
+    _TORCHVISION_MNIST_AVAILABLE,
+    cli_lightning_logo,
+)
+
+if _TORCHVISION_AVAILABLE and _TORCHVISION_MNIST_AVAILABLE:
     from torchvision import transforms
     from torchvision.datasets.mnist import MNIST
 else:

pl_examples/basic_examples/mnist_datamodule.py

@@ -17,10 +17,10 @@
 
 from torch.utils.data import DataLoader, random_split
 
-from pl_examples import _DATASETS_PATH, _TORCHVISION_AVAILABLE
+from pl_examples import _DATASETS_PATH, _TORCHVISION_AVAILABLE, _TORCHVISION_MNIST_AVAILABLE
 from pytorch_lightning import LightningDataModule
 
-if _TORCHVISION_AVAILABLE:
+if _TORCHVISION_AVAILABLE and _TORCHVISION_MNIST_AVAILABLE:
     from torchvision import transforms as transform_lib
     from torchvision.datasets import MNIST
 else:

pl_examples/domain_templates/generative_adversarial_net.py

@@ -26,15 +26,19 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F  # noqa
-import torchvision
-import torchvision.transforms as transforms
 from torch.utils.data import DataLoader
-from torchvision.datasets import MNIST
 
-from pl_examples import cli_lightning_logo
+from pl_examples import _TORCHVISION_AVAILABLE, _TORCHVISION_MNIST_AVAILABLE, cli_lightning_logo
 from pytorch_lightning.core import LightningDataModule, LightningModule
 from pytorch_lightning.trainer import Trainer
 
+if _TORCHVISION_AVAILABLE and _TORCHVISION_MNIST_AVAILABLE:
+    import torchvision
+    import torchvision.transforms as transforms
+    from torchvision.datasets import MNIST
+else:
+    from tests.helpers.datasets import MNIST
+
 
 class Generator(nn.Module):
     """