chore: Also add datapaths to job
This commit is contained in:
parent
81c767371e
commit
b74ae7083a
GPG key ID:
38297DE43F75FFE2
2 changed files with 86 additions and 45 deletions
|
|
@ -7,58 +7,95 @@ from torch.utils.data import DataLoader
|
|||
from dataset_loaders import dataset_called
|
||||
from trainers import OptunaTrainer, Trainer, FullTrainer
|
||||
|
||||
BATCH_SIZE = 64
|
||||
|
||||
if torch.accelerator.is_available():
|
||||
DEVICE = torch.accelerator.current_accelerator().type
|
||||
else:
|
||||
DEVICE = "cpu"
|
||||
def parse_arguments():
|
||||
parser = ArgumentParser(prog="NeuralCompression")
|
||||
parser.add_argument("--verbose", "-v", action="store_true", required=False,
|
||||
help="Enable verbose mode")
|
||||
|
||||
# hyper parameters
|
||||
context_length = 128
|
||||
dataparser = ArgumentParser(add_help=False)
|
||||
dataparser.add_argument("--data-root", type=str, required=False)
|
||||
dataparser.add_argument("--dataset", choices=dataset_called.keys(), required=True)
|
||||
|
||||
print(f"Running on device: {DEVICE}...")
|
||||
parser = ArgumentParser()
|
||||
parser.add_argument("--method", choices=["optuna", "train"], required=True)
|
||||
parser.add_argument("--model-path", type=str, required=False)
|
||||
modelparser = ArgumentParser(add_help=False)
|
||||
modelparser.add_argument("--model-path", type=str, required=True,
|
||||
help="Path to the model to load/save")
|
||||
|
||||
parser.add_argument_group("Data", "Data files or dataset to use")
|
||||
parser.add_argument("--data-root", type=str, required=False)
|
||||
parser.add_argument("dataset")
|
||||
args = parser.parse_args()
|
||||
fileparser = ArgumentParser(add_help=False)
|
||||
fileparser.add_argument("--input-file", "-i", required=False, type=str)
|
||||
fileparser.add_argument("--output-file", "-o", required=False, type=str)
|
||||
|
||||
print("Loading in the dataset...")
|
||||
if args.dataset in dataset_called:
|
||||
dataset = dataset_called[args.dataset](root=args.data_root, transform=lambda x: x.to(DEVICE))
|
||||
else:
|
||||
# TODO Allow to import arbitrary files
|
||||
raise NotImplementedError(f"Importing external datasets is not implemented yet")
|
||||
subparsers = parser.add_subparsers(dest="mode", required=True,
|
||||
help="Mode to run in")
|
||||
|
||||
dataset_length = len(dataset)
|
||||
print(f"Dataset size = {dataset_length}")
|
||||
# TODO
|
||||
fetch_parser = subparsers.add_parser("fetch", parents=[dataparser],
|
||||
help="Only fetch the dataset, then exit")
|
||||
|
||||
training_size = ceil(0.8 * dataset_length)
|
||||
train_parser = subparsers.add_parser("train", parents=[dataparser, modelparser])
|
||||
|
||||
print(f"Training set size = {training_size}, Validation set size {dataset_length - training_size}")
|
||||
# TODO
|
||||
compress_parser = subparsers.add_parser("compress", parents=[modelparser, fileparser])
|
||||
|
||||
train_set, validate_set = torch.utils.data.random_split(dataset,
|
||||
[training_size, dataset_length - training_size])
|
||||
training_loader = DataLoader(train_set, batch_size=BATCH_SIZE, shuffle=True)
|
||||
validation_loader = DataLoader(validate_set, batch_size=BATCH_SIZE, shuffle=False)
|
||||
loss_fn = torch.nn.CrossEntropyLoss()
|
||||
# TODO
|
||||
decompress_parser = subparsers.add_parser("decompress", parents=[modelparser, fileparser])
|
||||
|
||||
model = None
|
||||
if args.model_path is not None:
|
||||
print("Loading the model...")
|
||||
model = torch.load(args.model_path)
|
||||
return parser.parse_args()
|
||||
|
||||
trainer: Trainer = OptunaTrainer() if args.method == "optuna" else FullTrainer()
|
||||
|
||||
trainer.execute(
|
||||
model=model,
|
||||
train_loader=training_loader,
|
||||
validation_loader=validation_loader,
|
||||
loss_fn=loss_fn,
|
||||
n_epochs=200,
|
||||
device=DEVICE
|
||||
)
|
||||
def main():
|
||||
BATCH_SIZE = 64
|
||||
|
||||
# hyper parameters
|
||||
context_length = 128
|
||||
|
||||
args = parse_arguments()
|
||||
|
||||
if torch.accelerator.is_available():
|
||||
DEVICE = torch.accelerator.current_accelerator().type
|
||||
else:
|
||||
DEVICE = "cpu"
|
||||
print(f"Running on device: {DEVICE}...")
|
||||
|
||||
print("Loading in the dataset...")
|
||||
if args.dataset in dataset_called:
|
||||
dataset = dataset_called[args.dataset](root=args.data_root, transform=lambda x: x.to(DEVICE))
|
||||
else:
|
||||
# TODO Allow to import arbitrary files
|
||||
raise NotImplementedError(f"Importing external datasets is not implemented yet")
|
||||
|
||||
if args.mode == 'fetch':
|
||||
# TODO More to earlier in chain, because now everything is converted into tensors as well?
|
||||
exit(0)
|
||||
|
||||
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}")
|
||||
|
||||
train_set, validate_set = torch.utils.data.random_split(dataset, [training_size, dataset_length - training_size])
|
||||
training_loader = DataLoader(train_set, batch_size=BATCH_SIZE, shuffle=True)
|
||||
validation_loader = DataLoader(validate_set, batch_size=BATCH_SIZE, shuffle=False)
|
||||
loss_fn = torch.nn.CrossEntropyLoss()
|
||||
|
||||
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()
|
||||
|
||||
trainer.execute(
|
||||
model=model,
|
||||
train_loader=training_loader,
|
||||
validation_loader=validation_loader,
|
||||
loss_fn=loss_fn,
|
||||
n_epochs=200,
|
||||
device=DEVICE
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
|
|
|||
Reference in a new issue