tdpeuter/2025ML-project-neural_compression
tdpeuter/2025ML-project-neural_compression
Archived
1
Fork 0
Code Issues Pull requests Releases 1 Activity Actions

fix: enwik dataset fix

Browse source
This commit is contained in:
Robin Meersman 2025-11-28 09:27:37 +01:00
parent fe207962de
commit 0577eee601
2 changed files with 34 additions and 12 deletions
Download patch file Download diff file Expand all files Collapse all files

6
CNN-model/main_cnn.py
View file

@ -20,6 +20,7 @@ if __name__ == "__main__":
parser.add_argument("--model-path", type=str, required=False)
args = parser.parse_args()
print("Loading in the dataset...")
if args.method == "train":
dataset: Dataset = EnWik9DataSet(transform=lambda x: x.to(DEVICE))
elif args.method == "optuna":
@ -28,9 +29,11 @@ if __name__ == "__main__":
raise ValueError(f"Unknown method: {args.method}")
dataset_length = len(dataset)
print(f"Dataset size = {dataset_length}")
training_size = ceil(0.8 * dataset_length)
print(f"training set size = {training_size}, validation set size {dataset_length - training_size}")
print(f"Training set size = {training_size}, Validation set size {dataset_length - training_size}")
train_set, validate_set = torch.utils.data.random_split(dataset,
[training_size, dataset_length - training_size])
@ -40,6 +43,7 @@ if __name__ == "__main__":
model = None
if args.model_path is not None:
print("Loading the model...")
model = torch.load(args.model_path)
trainer: Trainer = OptunaTrainer() if args.method == "optuna" else FullTrainer()