90 lines
3.6 KiB
Python
90 lines
3.6 KiB
Python
# Imports
|
|
# =============================================================================
|
|
import os
|
|
from distutils.util import strtobool
|
|
from pathlib import Path
|
|
from argparse import ArgumentParser
|
|
|
|
import warnings
|
|
|
|
from pytorch_lightning import Trainer
|
|
from torch.utils.data import DataLoader
|
|
|
|
from dataset.dataset import TrajPairData
|
|
from lib.utils.config import Config
|
|
from lib.utils.logging import Logger
|
|
|
|
warnings.filterwarnings('ignore', category=FutureWarning)
|
|
warnings.filterwarnings('ignore', category=UserWarning)
|
|
|
|
_ROOT = Path(__file__).parent
|
|
|
|
# Paramter Configuration
|
|
# =============================================================================
|
|
# Argument Parser
|
|
main_arg_parser = ArgumentParser(description="parser for fast-neural-style")
|
|
|
|
# Main Parameters
|
|
main_arg_parser.add_argument("--main_debug", type=strtobool, default=False, help="")
|
|
main_arg_parser.add_argument("--main_eval", type=strtobool, default=False, help="")
|
|
main_arg_parser.add_argument("--main_seed", type=int, default=69, help="")
|
|
|
|
# Data Parameters
|
|
main_arg_parser.add_argument("--data_worker", type=int, default=10, help="")
|
|
main_arg_parser.add_argument("--data_batchsize", type=int, default=100, help="")
|
|
main_arg_parser.add_argument("--data_root", type=str, default='/data/', help="")
|
|
main_arg_parser.add_argument("--map_root", type=str, default='/res/maps', help="")
|
|
|
|
# Transformations
|
|
main_arg_parser.add_argument("--transformations_to_tensor", type=strtobool, default=False, help="")
|
|
|
|
# Transformations
|
|
main_arg_parser.add_argument("--train_outpath", type=str, default="output", help="")
|
|
main_arg_parser.add_argument("--train_version", type=strtobool, required=False, help="")
|
|
main_arg_parser.add_argument("--train_epochs", type=int, default=10, help="")
|
|
main_arg_parser.add_argument("--train_batch_size", type=int, default=512, help="")
|
|
main_arg_parser.add_argument("--train_lr", type=float, default=0.002, help="")
|
|
|
|
# Model
|
|
main_arg_parser.add_argument("--model_type", type=str, default="LeNetAE", help="")
|
|
main_arg_parser.add_argument("--model_activation", type=str, default="relu", help="")
|
|
main_arg_parser.add_argument("--model_filters", type=str, default="[32, 16, 4]", help="")
|
|
main_arg_parser.add_argument("--model_use_bias", type=strtobool, default=True, help="")
|
|
main_arg_parser.add_argument("--model_use_norm", type=strtobool, default=True, help="")
|
|
main_arg_parser.add_argument("--model_dropout", type=float, default=0.00, help="")
|
|
|
|
# Project
|
|
main_arg_parser.add_argument("--project_name", type=str, default='traj-gen', help="")
|
|
main_arg_parser.add_argument("--project_owner", type=str, default='si11ium', help="")
|
|
main_arg_parser.add_argument("--project_neptune_key", type=str, default=os.getenv('NEPTUNE_KEY'), help="")
|
|
|
|
# Parse it
|
|
args = main_arg_parser.parse_args()
|
|
config = Config.read_namespace(args)
|
|
|
|
################
|
|
# TESTING ONLY #
|
|
# =============================================================================
|
|
hparams = config.model_paramters
|
|
dataset = TrajPairData('data', mapname='tate', alternatives=10000, trajectories=2500)
|
|
dataloader = DataLoader(dataset=dataset.train_dataset, shuffle=True,
|
|
batch_size=hparams.data_param.batchsize,
|
|
num_workers=hparams.data_param.worker)
|
|
|
|
# Logger
|
|
# =============================================================================
|
|
logger = Logger(config, debug=True)
|
|
|
|
# Trainer
|
|
# =============================================================================
|
|
trainer = Trainer(logger=logger)
|
|
|
|
# Model
|
|
# =============================================================================
|
|
model = None
|
|
|
|
|
|
if __name__ == '__main__':
|
|
next(iter(dataloader))
|
|
pass
|