#! /usr/bin/env python
# coding: utf-8
import os, sys
import numpy as np
import cv2 as cv
import mxnet as mx
import matplotlib.pyplot as plt
import glob
import mxnet.metric
import warnings
from logger.logger_v4 import Log
logging = Log()
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from dataset.image_folder_dataset import ImageFolderDataIter, ImageFolderDataset, ImageFolderDataLoader
from utils.file_func import get_class_name, get_function_name
from model_zoo.mobilenet_v2 import MobileNetV20Transform
class Trainer():
def __init__(self, config):
self.config = config
self.ctxs = [mx.gpu(int(i)) for i in config.gpus.split(',')]
self.net = self.load_model_zoo(pretrained=self.config.pretrained_name,
epoch=self.config.pretrained_epoch)
self.net.collect_params().reset_ctx(self.ctxs)
self.trainer = mx.gluon.Trainer(self.net.collect_params(), 'sgd', {
def __init__(self, name='default', logger=Log()):
self.name = name
self.logging = logger