https://github.com/wyhwong/classifier-trains

This repository aims to provide a primitive tool to finetune state-of-the-art models with PyTorch implementation, similar to Nvidia TAO but with more flexibility in augmentation and models.
https://github.com/wyhwong/classifier-trains

computer-vision deep-learning image-classifier pytorch pytorch-lightning

Last synced: 6 months ago
JSON representation

This repository aims to provide a primitive tool to finetune state-of-the-art models with PyTorch implementation, similar to Nvidia TAO but with more flexibility in augmentation and models.

• Host: GitHub
• URL: https://github.com/wyhwong/classifier-trains
• Owner: wyhwong
• License: mit
• Created: 2023-01-21T03:19:07.000Z (over 3 years ago)
• Default Branch: main
• Last Pushed: 2025-12-31T16:10:12.000Z (6 months ago)
• Last Synced: 2026-01-05T01:03:35.050Z (6 months ago)
• Topics: computer-vision, deep-learning, image-classifier, pytorch, pytorch-lightning
• Language: Python
• Homepage: https://pypi.org/project/classifier_trains/
• Size: 3.84 MB
• Stars: 3
• Watchers: 1
• Forks: 2
• Open Issues: 0
• Metadata Files:
  • Readme: README.md
  • License: LICENSE
  • Codeowners: .github/CODEOWNERS

Awesome Lists containing this project

README

          
# Classifier-trains

[![PyPI version](https://badge.fury.io/py/classifier-trains.svg)](https://pypi.org/project/classifier-trains)

[![Python version](https://img.shields.io/pypi/pyversions/classifier-trains)](https://pypi.org/project/classifier-trains/)

[![Downloads](https://img.shields.io/pepy/dt/classifier-trains)](https://github.com/wyhwong/classifier-trains)

[![license](https://img.shields.io/badge/license-MIT-orange.svg)](https://github.com/wyhwong/classifier-trains/blob/main/LICENSE)

[![CI](https://github.com/wyhwong/classifier-trains/actions/workflows/main.yml/badge.svg)](https://github.com/wyhwong/classifier-trains/actions/workflows/main.yml/)

Classifier-trains is a package aims to provide a primitive tool to finetune state-of-the-art models with PyTorch implementation, similar to Nvidia TAO but with more flexibility in augmentation and models. Like TAO classification, all parameters are configurable in yaml file to minimize the need of modifying scripts.

Goal: Get config ready, get dataset ready, no coding (hopefully :D), start PyTorch training.

## Install as a package

```bash

pip3 install classifier-trains

```

## Usage

Please see [README.md](./src/README.md) inside `src` folder.

## UML Diagram

Please see [UML Diagram](./docs/README.md) for the class diagram.

## After Training

The logs and checkpoints will be saved in the output directory, and logs are in tensorboard format. In tensorboard, you will be able to see the ROC curve, sample images in training, parameters like learning rate and momentum, and metrics like accuracy and loss.

```bash

# Run tensorboard

tensorboard --logdir 

```

![ROC curve in tensorboard](./docs/preview-tensorboard-1.png)

![Sample images and parameters](./docs/preview-tensorboard-2.png)

## Profile Report

If you run the training with profiling, a profile report will be generated in the output directory. You can see the time spent on each function. It gives you a better understanding of the performance of the training process, and a sense of where to optimize.

![Profile report](./docs/preview-profile-report.png)

## Features To Be Developed

1. Implement detector training

## Author

[@wyhwong](https://github.com/wyhwong)