A clear PyTorch template for swift model building.
- Well-organized project template out of the box.
- Automatically record the model version (by saving the git commit hash) for later reproduction.
- Automatically start TensorBoard for you.
- Use JSON file or command line arguments to specify arguments.
- The results of each experiment are properly stored.
- Modifiy the model structures
models/build.py. - Update the loss functions used
solver/loss.py. - Update the data loading process
data/dataset.py&data/loader.py. - Add metrics that can measure your model's performance
metrics/eval.py. - Update sampling functions & logging functions so you can see the results with TensorBoard
solver/solver.py! - Add a shell script that run your model
scripts/{exp_id}-model_key_config.sh. - Start training, evaluating or inferencing by running the above script!
+--- .gitignore
+--- archive (generated files & dataset)
| +--- README.md
+--- bin (utility script)
| +--- README.md
| +--- template.py
+--- config.py (options)
+--- data (data fetching related)
| +--- dataset.py
| +--- fetcher.py
| +--- loader.py
| +--- README.md
+--- expr (experiment directory)
+--- main.py (everything start from here)
+--- metrics (metric used)
| +--- eval.py
| +--- fid.py
| +--- README.md
+--- models (model architecture related)
| +--- build.py (the wrapper for models)
| +--- discriminator.py
| +--- generator.py
| +--- layers.py
| +--- mapping_network.py
| +--- README.md
+--- README.md
+--- requirements.txt
+--- scripts (training related shell scripts)
| +--- train.sh
+--- solver (training related)
| +--- loss.py
| +--- misc.py
| +--- solver.py
| +--- utils.py
+--- utils (utility functions)
| +--- checkpoint.py
| +--- file.py
| +--- image.py
| +--- logger.py
| +--- misc.py
| +--- model.py
I referred StarGAN v2's official implementation when crafting this template, so don't be surprised if you find some code is similar.