Skip to content
Snippets Groups Projects
train.py 628 B
Newer Older
Moi's avatar
Moi committed
#! /usr/bin/env python
Moi's avatar
Moi committed
from torch.multiprocessing import set_start_method
Moi's avatar
Moi committed
from src.options import TrainingOptions
from src.data import CustomDatasetDataLoader
from src.models import create_model, train_model

if __name__ == "__main__":
Moi's avatar
Moi committed
    set_start_method('spawn')
Moi's avatar
Moi committed
    options = TrainingOptions()
    options.parse()
Moi's avatar
Moi committed

Moi's avatar
Moi committed
    dataset = CustomDatasetDataLoader(options.options)
Moi's avatar
Moi committed

    model = create_model(options)
Moi's avatar
Moi committed
    if options.options.verbose:
Moi's avatar
Moi committed
        print("-----------------Networks---------------")
        print(model)
        print("---------------End networks-------------")
Moi's avatar
Moi committed

Moi's avatar
Moi committed
    train_model(model, dataset, options.options)