Tweaked dataloader for more reliable loading

This commit is contained in:
2024-05-13 15:38:50 +01:00
parent 8204109561
commit 223b064564
2 changed files with 7 additions and 1 deletions
+2
View File
@@ -160,3 +160,5 @@ cython_debug/
# and can be added to the global gitignore or merged into this file. For a more nuclear
# option (not recommended) you can uncomment the following to ignore the entire idea folder.
#.idea/
datasets/
+5 -1
View File
@@ -1,10 +1,14 @@
from pathlib import Path
from torchvision.datasets import Caltech256
from torch.utils.data import random_split
from torch.utils.data import BatchSampler
PROJECT_ROOT = Path(__file__).parent.parent
def get_dataset(split: (float, float, float) = (0.7, 0.1, 0.2), *args, **kwargs):
ds = Caltech256("../datasets/", download=True)
ds = Caltech256(PROJECT_ROOT / "datasets/", download=True)
train, test, val = (
BatchSampler(i, *args, **kwargs) for i in random_split(ds, split)
)