Yoga-Pose-Classification/src/train.py

from config import MODELS_DIR, IMAGES
import matplotlib.pyplot as plt
import torch
from fastai.vision.all import *
from fastai.metrics import error_rate, accuracy

torch.cuda.set_device(0)
import mlflow
import os


mlflow.set_tracking_uri(os.getenv("MLFLOW_TRACKING_URI"))


def get_experiment_id(name):
    exp = mlflow.get_experiment_by_name(name)
    if exp is None:
        exp_id = mlflow.create_experiment(name)
        return exp_id
    return exp.experiment_id


exp_id = get_experiment_id("yoga_lover")

if __name__ == "__main__":



    data = ImageDataLoaders.from_folder(IMAGES, valid_pct=0.2, item_tfms=Resize(224))

    print("Training the model...")
    learn = vision_learner(data, resnet34, metrics=[accuracy, error_rate])
    mlflow.fastai.autolog()
    learn.remove_cb(ProgressCallback)
    with mlflow.start_run(experiment_id=exp_id):
        learn.fine_tune(3)
    print("Training completed.")
    model_path="/content/Yoga-Pose-Classification/Saved_Models/"
    os.mkdir(model_path)
    print("Saving the model...")
    
    learn.save(model_path+MODELS_DIR, with_opt=False)
    print("done.")