class_docstr: 🤗 Constructor for the image classifier trainer
code: def __init__(
      self,
      id2label,
      label2id,
      model_name: str,
      train: torch.utils.data.Dataset,
      test: torch.utils.data.Dataset,
      output_dir="./models",
      max_epochs=1,
      workers=mp.cpu_count(),
      batch_size=8,
      lr=2e-5,
      weight_decay=0.01,
      momentum=0.9,
      gamma=0.96,
      pretrained=True,
      force_cpu=False,
      requires_grad=False,  # True: Only last layer of the classifier are updated
      load_best_model_at_end=True,
      classification_report_digits=4,
      parallelized=False,
      lite=False,
  ):
      self.train = train
      self.test = test
      self.output_dir = (
          output_dir if output_dir.endswith("/") else output_dir + "/"
      )
      self.lr = lr
      self.weight_decay = weight_decay
      self.momentum = momentum
      self.gamma = gamma
      self.batch_size = batch_size
      self.max_epochs = max_epochs
      self.workers = workers
      self.model_name = model_name
      self.pretrained = pretrained
      self.ids2labels = id2label
      self.labels2ids = label2id
      self.best_acc = 0
      self.best_path = ""
      self.logs_path = self.output_dir + "logs/"
      self.config_path = self.output_dir + "config.json"
      self.current_date = datetime.today().strftime("%Y-%m-%d-%H-%M-%S")
      self.config = {}
      self.requires_grad = requires_grad
      self.load_best_model_at_end = load_best_model_at_end
      self.classification_report_digits = classification_report_digits
      self.parallelized = parallelized
      self.lite = lite
  
      self.tensor_board = SummaryWriter()
  
      self.data_loader_train = torch.utils.data.DataLoader(
          self.train,
          batch_size=self.batch_size,
          shuffle=True,
          num_workers=self.workers,
      )
      self.data_loader_test = torch.utils.data.DataLoader(
          self.test,
          batch_size=self.batch_size,
          shuffle=True,
          num_workers=self.workers,
      )
  
      self.image_example = next(iter(self.data_loader_train))[0]
  
      print(self.ids2labels)
      print(self.labels2ids)
  
      # Processing device (CPU / GPU)
      if force_cpu == True:
          self.device = "cpu"
      else:
          self.device = torch.device(
              "cuda:0" if torch.cuda.is_available() else "cpu"
          )
  
      # Setup the metric
      self.metric = torchmetrics.Accuracy(task="multiclass")
  
      self.num_classes = len(self.ids2labels.keys())
      self.config["num_classes"] = self.num_classes
  
      # Open the logs files
      self.__openLogs()
      self.__openLogsLossAcc()
  
      # Load the model from the TorchVision Models Zoo
      if pretrained:
          print("Load the pre-trained model " + self.model_name)
          self.model = models.__dict__[self.model_name](pretrained=True)
      else:
          print("Load the model " + self.model_name)
          self.model = models.__dict__[self.model_name]()
  
      print(self.model)
  
      # """
      # 🏗️ Build the model
      # """
      if self.model_name.startswith(
          "alexnet"
      ) or self.model_name.startswith("vgg"):
          self.model.classifier[6] = nn.Linear(
              self.model.classifier[6].in_features, self.num_classes
          )
  
          if self.parallelized == True:
              self.model.features = torch.nn.DataParallel(
                  self.model.features
              )
  
          self.model.to(self.device)
          self.config["hidden_size"] = self.model.classifier[
              6
          ].in_features
      else:
          if (
              self.model_name.startswith("resnet")
              or self.model_name.startswith("googlenet")
              or self.model_name.startswith("resnext")
              or self.model_name.startswith("shufflenet_v2")
              or self.model_name.startswith("wide_resnet")
          ):
              self.model.fc = nn.Linear(
                  self.model.fc.in_features, self.num_classes
              )
              self.config["hidden_size"] = self.model.fc.in_features
  
          elif self.model_name.startswith("squeezenet"):
              self.model.classifier[1] = nn.Conv2d(
                  512,
                  self.num_classes,
                  kernel_size=(1, 1),
                  stride=(1, 1),
              )
              self.config["hidden_size"] = 512
  
          elif self.model_name.startswith("densenet"):
              self.model.classifier = nn.Linear(
                  self.model.classifier.in_features, self.num_classes
              )
              self.config[
                  "hidden_size"
              ] = self.model.classifier.in_features
  
          elif self.model_name.startswith("inception_v3"):
              self.model.AuxLogits.fc = nn.Linear(
                  self.model.AuxLogits.fc.in_features, self.num_classes
              )
              self.model.fc = nn.Linear(2048, self.num_classes)
              self.config[
                  "hidden_size"
              ] = self.model.AuxLogits.fc.in_features
  
          elif self.model_name.startswith(
              "mobilenet_v2"
          ) or self.model_name.startswith("mnasnet"):
              self.model.classifier[1] = nn.Linear(
                  self.model.classifier[1].in_features, self.num_classes
              )
              self.config["hidden_size"] = self.model.classifier[
                  1
              ].in_features
  
          elif self.model_name.startswith("mobilenet_v3"):
              self.model.classifier[3] = nn.Linear(
                  self.model.classifier[3].in_features, self.num_classes
              )
              self.config["hidden_size"] = self.model.classifier[
                  3
              ].in_features
  
          if self.parallelized == True:
              self.model = torch.nn.DataParallel(self.model)
  
          self.model = self.model.to(self.device)
  
      if self.requires_grad:
          for param in self.model.parameters():
              param.requires_grad = False
  
      archi = ""
      archi += "=" * 50 + "\n"
      archi += "Model architecture:" + "\n"
      archi += "=" * 50 + "\n"
      archi += str(self.model) + "\n"
      archi += "=" * 50 + "\n"
      print(archi)
  
      # Write in logs
      self.__openLogs()
      self.logs_file.write(archi + "\n")
      self.logs_file.close()
  
      self.criterion = nn.CrossEntropyLoss().to(self.device)
      self.optimizer = torch.optim.Adam(
          self.model.parameters(), self.lr
      )
      # self.optimizer = torch.optim.SGD(self.model.parameters(), self.lr, momentum=self.momentum, weight_decay=self.weight_decay)
      self.scheduler = ExponentialLR(self.optimizer, gamma=self.gamma)
  
      self.config["id2label"] = self.ids2labels
      self.config["label2id"] = self.labels2ids
      self.config["architectures"] = [self.model_name]
  
      with open(self.config_path, "w") as json_file:
          json.dump(self.config, json_file, indent=4)
  
      """
      ⚙️ Train the given model on the dataset
      """
      print("Start Training!")
      self.training()
  
      if self.load_best_model_at_end:
          print("\033[95mBest model loaded!\033[0m")
          self.model = torch.load(
              self.output_dir + "best_model.pth",
              map_location=self.device,
          )
  
      self.evaluate_f1_score()
  
      # Close the logs file
      self.logs_file.close()
      self.logs_loss_train.close()
      self.logs_acc_train.close()
  
      if self.lite == True:
          lite_path = output_dir + "model.ptl"
          print(lite_path)
          traced_script_module = torch.jit.trace(
              self.model.to("cpu"), self.image_example.to("cpu")
          )
          traced_script_module_optimized = optimize_for_mobile(
              traced_script_module
          )
          traced_script_module_optimized._save_for_lite_interpreter(
              lite_path
          )
          print("Lite exported!")
  
imports: import os
  import json
  import math
  import random
  import argparse
  from pathlib import Path
  import multiprocessing
  from datetime import datetime
  from collections import Counter
  import torch
  import torchmetrics
  import torch.nn
  from torchmetrics import Accuracy
  import torchvision.models
  from torch.utils.data import DataLoader
  from torch.utils.tensorboard import SummaryWriter
  from torch.optim.lr_scheduler import ExponentialLR
  from torch.utils.mobile_optimizer import optimize_for_mobile
  from sklearn import metrics
  from sklearn.metrics import accuracy_score
  from sklearn.metrics import classification_report
  from tqdm import tqdm
  from tabulate import tabulate
  from PIL import Image, ImageEnhance
label_desc: Use image node_id 1 to create a new TorchVisionClassifierTrainer object  with example: obj = TorchVisionClassifierTrainer(id2label, label2id, model_name, train, test, output_dir, max_epochs, workers, batch_size, lr, weight_decay, momentum, gamma, pretrained, force_cpu, requires_grad, load_best_model_at_end, classification_report_digits, parallelized, lite)