[FIX] Use ImageNet train and add to __init__

fsschneider · Oct 21, 2021 · d9ba931 · d9ba931
1 parent 52bed6b
commit d9ba931
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/deepobs/pytorch/datasets/imagenet.py b/deepobs/pytorch/datasets/imagenet.py
@@ -89,12 +89,12 @@ def _make_train_and_valid_dataloader(self):
 
         train_dataset = datasets.ImageNet(
             root=os.path.join(config.get_data_dir(), "imagenet/pytorch"),
-            split="val",  # TODO Change to train again
+            split="train",
             transform=transform,
         )
         valid_dataset = datasets.ImageNet(
             root=os.path.join(config.get_data_dir(), "imagenet/pytorch"),
-            split="val",  # TODO Change to train again
+            split="train",
             transform=transform,
         )
         train_loader, valid_loader = self._make_train_and_valid_dataloader_helper(

diff --git a/deepobs/pytorch/testproblems/__init__.py b/deepobs/pytorch/testproblems/__init__.py
@@ -13,6 +13,8 @@
 from .fmnist_logreg import fmnist_logreg
 from .fmnist_mlp import fmnist_mlp
 from .fmnist_vae import fmnist_vae
+from .imagenet_vgg16 import imagenet_vgg16
+from .imagenet_vgg19 import imagenet_vgg19
 from .mnist_2c2d import mnist_2c2d
 from .mnist_logreg import mnist_logreg
 from .mnist_mlp import mnist_mlp