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

feat: Size arg

Browse source
This commit is contained in:
Tibo De Peuter 2025-12-10 16:29:39 +01:00
parent 6c5908e6ae
commit 67310ce4f4
Signed by: tdpeuter
GPG key ID: 38297DE43F75FFE2
2 changed files with 6 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

5
main.py
View file

@ -12,11 +12,14 @@ def main():
match args.mode: match args.mode:
case 'train': case 'train':
size = None size = args.size
if args.method == 'optuna': if args.method == 'optuna':
size = 2 ** 12 size = 2 ** 12
print(f"Using size {size} for optuna (was {args.size})")
if args.debug: if args.debug:
size = 2 ** 10 size = 2 ** 10
print(f"Using size {size} for debug (was {args.size})")
train( train(
device=device, device=device,
dataset=args.dataset, dataset=args.dataset,

2
src/args.py
View file

@ -35,6 +35,8 @@ def parse_arguments():
train_parser.add_argument("--method", train_parser.add_argument("--method",
choices=["fetch", "optuna", "full"], required=True, choices=["fetch", "optuna", "full"], required=True,
help="Method to use for training") help="Method to use for training")
train_parser.add_argument("--size", "-s", type=int, required=False,
help="Size of the subset of the dataset to use")
compress_parser = subparsers.add_parser("compress", parents=[modelparser, fileparser]) compress_parser = subparsers.add_parser("compress", parents=[modelparser, fileparser])