UniTrain

UniTrain is an open-source, unified platform for effortless machine learning model training, evaluation, and deployment across diverse tasks. Experience seamless experimentation and model deployment with UniTrain.

Installation instruction

Install the UniTrain module using:

pip install -i https://test.pypi.org/simple/ UniTrain==0.2.3

Install torch library using:

pip install torch

Note : You can use your local device or any online notebooks like Google Colab or Kagle for training the models, as explained below.
However,using 'Google Colab' would be preferred because of its simple user-friendly interface and the computing power that it brings with itself.

Note: For Google Colab use "!" before every command.

Training

Classification

• Create a "data" folder within a "content" folder.
• The "data" folder will contain three different folders named "train", "test" and "eval" used for training, testing, and evaluation purposes.
• Each of the "train", "test" and "eval" folders contains data sets of different categories on which you want to use your model
• Run the following code to train your model and you can change the default arguments with your custom arguments

import UniTrain
from UniTrain.utils.classification import get_data_loader, train_model
from UniTrain.models.classification import ResNet9
from UniTrain.utils.classification import parse_folder
import torch

if parse_folder("/content/data/"):
  train_dataloader = get_data_loader("/content/data/", 32, True, split='train')
  test_dataloader = get_data_loader("/content/data/", 32, True, split='test')

  model = ResNet9(num_classes=6)
  model.to(torch.device('cuda'))

  train_model(model, train_dataloader, test_dataloader,
              num_epochs=10, learning_rate=1e-3, checkpoint_dir='checkpoints',logger = "training.log", device=torch.device('cuda'))

Segmentation

• Create a "data" folder within a "content" folder.
• The "data" folder will contain three different folders named "train", "test" and "eval" used for training, testing, and evaluation purposes.
• Each of the "train", "test" and "eval" folders contains data sets of different categories on which you want to use your model
• Run the following code to train your model and you can change the default arguments with your custom arguments

import UniTrain
from UniTrain.utils.classification import get_data_loader, train_model
from UniTrain.models.classification import ResNet9
from UniTrain.utils.classification import parse_folder
import torch

if parse_folder("/content/data/"):
  train_dataloader = get_data_loader("/content/data/", 32, True, split='train')
  test_dataloader = get_data_loader("/content/data/", 32, True, split='test')

  model = ResNet9(num_classes=6)
  model.to(torch.device('cuda'))

  train_model(model, train_dataloader, test_dataloader,
              num_epochs=10, learning_rate=1e-3, checkpoint_dir='checkpoints',logger = "training.log", device=torch.device('cuda'))

DCGAN

Adding Data for Training

• For the data create a folder 'data'->'realimages'->'images'-> Add your data here
• No need to create fake images as it would be generated by the untrained generator

Training the model ( both dicriminator and generator )

• Run the following code to train your model and you can change the default arguments with your custom arguments

import UniTrain
from UniTrain.utils.DCGAN import parse_folder, get_data_loader, train_model
from UniTrain.models.DCGAN import disc , gen
import glob
import torch

def main():
  result = parse_folder('data')
  if result:

    real_image_paths = glob.glob("data/real_images/images/*.jpg")
    real_image_paths = glob.glob("data/real_images/images/*.png")
    real_image_paths = glob.glob("data/real_images/images/*.jpeg")
    
    train_dataloader = get_data_loader('data', 128 )
    discriminator_model = disc.discriminator
    generator_model = gen.generator
    train_model( discriminator_model, generator_model, train_dataloader, batch_size = 128 ,  num_epochs = 25, learning_rate = 1e-3, torch.device('cpu'),checkpoint_dir='checkpoints')

  else:
    print("Invalid dataset folder.")
    return None

if __name__ == '__main__':
    main()

`