cifar100-enriched / prepare.py
MarkusStoll's picture
dataset in dockerfile
1750280
import pickle
import datasets
import os
if __name__ == "__main__":
cache_file = "dataset_cache.pkl"
if os.path.exists(cache_file):
# Load dataset from cache
with open(cache_file, "rb") as file:
dataset = pickle.load(file)
print("Dataset loaded from cache.")
else:
# Load dataset using datasets.load_dataset()
dataset = datasets.load_dataset("renumics/cifar100-enriched", split="train")
print("Dataset loaded using datasets.load_dataset().")
# Save dataset to cache
with open(cache_file, "wb") as file:
pickle.dump(dataset, file)
print("Dataset saved to cache.")