beta0.1

app.py

@@ -0,0 +1,35 @@
+import fire
+
+from loguru import logger
+from utils import project_manager
+from utils import cache_data
+from utils import load_cache
+
+
+class App:
+
+    def __init__(self):
+        logger.info("\nHello baby~")
+
+    def create(self, project_name: str, single: bool = False):
+        logger.info("\nCreate Project ----> {}".format(project_name))
+        pm = project_manager.ProjectManager()
+        pm.create_project(project_name, single)
+
+    def cache(self, project_name: str, base_path: str, search_type: str = "name"):
+        logger.info("\nCaching Data ----> {}\nPath ----> {}".format(project_name, base_path))
+        cache = cache_data.CacheData(project_name)
+        cache.cache(base_path, search_type)
+        pass
+
+    def test_load(self):
+        load = load_cache.GetLoader("test1")
+        val = load.loaders['val']
+        val = iter(val)
+        for inputs, labels, labels_length in val:
+            print(inputs, labels, labels_length)
+
+
+
+if __name__ == '__main__':
+    fire.Fire(App)

configs/__init__.py

@@ -0,0 +1,3 @@
+from .base import *
+
+

configs/base.py

@@ -0,0 +1,61 @@
+import os
+import json
+import yaml
+
+
+class Config(object):
+
+    def __init__(self, project_name):
+        self.project_name = project_name
+        self.base_path = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "projects")
+        self.config_dict = {
+            "System": {
+                "Project": None,
+                "GPU": True,
+                "GPU_ID": 0,
+                "Allow_Ext": ["jpg", "jpeg", "png", "bmp"],
+                "Path": "",
+                "Val": 0.03
+            },
+            "Model": {
+                "ImageWidth": -1,
+                "ImageHeight": 64,
+                "ImageChannel": 1,
+                "CharSet": [],
+                "Word": False
+            },
+            "Train": {
+                "BATCH_SIZE": 32,
+                "TEST_BATCH_SIZE": 32,
+                'CNN': {
+                    "NAME": "ddddOcr",
+                },
+                'DROPOUT': 0.3,
+                'OPTIMIZER': 'SGD',
+                "TEST_STEP": 1000,
+                "TARGET": {
+                    "Accuracy": 0.97,
+                    "Epoch": 200,
+                    "Cost": 0.005
+                },
+                "LR": 0.01
+            }
+        }
+
+
+    def make_config(self, config_dict=None, single: bool = False):
+        if not config_dict:
+            config_dict = self.config_dict
+            if single:
+                config_dict['Model']['Word'] = True
+
+            config_dict["System"]["Project"] = self.project_name
+        config_path = os.path.join(self.base_path, self.project_name, "config.yaml")
+        with open(config_path, 'w', encoding="utf-8") as f:
+            yaml.dump(config_dict, f, allow_unicode=True, default_flow_style=None, indent=4)
+
+    def load_config(self):
+        config_path = os.path.join(self.base_path, self.project_name, "config.yaml")
+        with open(config_path, 'r', encoding="utf-8") as f:
+            config_dict = yaml.load(f, Loader=yaml.FullLoader)
+        return config_dict

nets/__init__.py

@@ -0,0 +1,186 @@
+import json
+
+from .backbone import *
+import torch
+
+
+class Net(torch.nn.Module):
+    def __init__(self, conf):
+        super(Net, self).__init__()
+        self.backbones_list = {
+            "ddddocr": DdddOcr,
+            "effnetv2_l": effnetv2_l,
+            "effnetv2_m": effnetv2_m,
+            "effnetv2_xl": effnetv2_xl,
+            "effnetv2_s": effnetv2_s,
+            "mobilenetv2": mobilenetv2,
+            "mobilenetv3_s": MobileNetV3_Small,
+            "mobilenetv3_l": MobileNetV3_Large
+        }
+
+        self.optimizers_list = {
+            "SGD": torch.optim.SGD,
+            "Adam": torch.optim.Adam,
+        }
+        self.conf = conf
+        self.image_channel = self.conf['Model']['ImageChannel']
+        self.resize = [int(self.conf['Model']['ImageWidth']), int(self.conf['Model']['ImageHeight'])]
+        self.charset = self.conf['Model']['CharSet']
+        self.charset_len = len(self.charset)
+        self.backbone = self.conf['Train']['CNN']['NAME']
+        self.paramters = []
+
+        if self.backbone in self.backbones_list:
+            test_cnn = self.backbones_list[self.backbone](nc=1)
+            x = torch.randn(1, 1, 64, 224)
+            test_features = test_cnn(x)
+            del x
+            del test_cnn
+            self.out_size = test_features.size()[1] * test_features.size()[2]
+            self.cnn = self.backbones_list[self.backbone](nc=self.image_channel)
+        else:
+            raise Exception("{} is not found in backbones! backbone list : {}".format(self.backbone, json.dumps(
+                list(self.backbones_list.keys()))))
+        self.paramters.append({'params': self.cnn.parameters()})
+
+        self.word = self.conf['Model']['Word']
+        if not self.word:
+            self.dropout = self.conf['Train']['DROPOUT']
+            self.lstm = torch.nn.LSTM(input_size=self.out_size, hidden_size=self.out_size, bidirectional=True,
+                                      dropout=self.dropout, num_layers=1)
+            self.paramters.append({'params': self.lstm.parameters()})
+
+            self.loss = torch.nn.CTCLoss(blank=0, reduction='mean')
+        else:
+            self.lstm = None
+            self.loss = torch.nn.CrossEntropyLoss()
+
+        self.paramters.append({'params': self.loss.parameters()})
+
+        self.fc = torch.nn.Linear(in_features=self.out_size * 2, out_features=self.charset_len)
+        self.paramters.append({'params': self.fc.parameters()})
+
+        self.lr = self.conf['Train']['LR']
+
+        self.optim = self.conf['Train']['OPTIMIZER']
+        if self.optim in self.optimizers_list:
+            if self.optim == "SGD":
+                self.optimizer = self.optimizers_list[self.optim](self.paramters, lr=self.lr, momentum=0.9)
+            else:
+                self.optimizer = self.optimizers_list[self.optim](self.paramters, lr=self.lr, betas=(0.9, 0.99))
+        else:
+            raise Exception("{} is not found in optimizers! optimizers list : {}".format(self.optim, json.dumps(
+                list(self.optimizers_list.keys()))))
+
+        self.scheduler = torch.optim.lr_scheduler.ExponentialLR(self.optimizer, gamma=0.98)
+
+    def forward(self, inputs):
+        predict = self.get_features(inputs)
+        if self.word:
+            outputs = predict.max(1)
+        else:
+            outputs = predict.max(2)[1].transpose(0, 1)
+        return outputs
+
+    def get_features(self, inputs):
+        outputs = self.cnn(inputs)
+        if not self.word:
+            outputs = outputs.permute(3, 0, 1, 2)
+            w, b, c, h = outputs.shape
+            outputs = outputs.view(w, b, c * h)
+            outputs, _ = self.lstm(outputs)
+            time_step, batch_size, h = outputs.shape
+            outputs = outputs.view(time_step * batch_size, h)
+            outputs = self.fc(outputs)
+            outputs = outputs.view(time_step, batch_size, -1)
+        else:
+            outputs = self.fc(outputs)
+        return outputs
+
+    def training(self, inputs, labels, labels_length):
+        outputs = self.get_features(inputs)
+        loss, lr = self.get_loss(outputs, labels, labels_length)
+        return loss, lr
+
+    def testing(self, inputs, labels, labels_length):
+        predict = self.get_feature(inputs)
+        pred_decode_labels = []
+        labels_list = []
+        correct_list = []
+        error_list = []
+        i = 0
+        labels = labels.tolist()
+        if self.word:
+            outputs = predict.max(1)[1]
+            for pred_labels in outputs:
+                pred_decode_labels.append(pred_labels)
+        else:
+            outputs = predict.max(2)[1].transpose(0, 1)
+            for pred_labels in outputs:
+                decoded = []
+                last_item = 0
+                for item in pred_labels:
+                    item = item.item()
+                    if item == last_item:
+                        continue
+                    else:
+                        last_item = item
+                    if item != 0:
+                        decoded.append(item)
+                pred_decode_labels.append(decoded)
+
+        for idx in labels_length.tolist():
+            labels_list.append(labels[i: i + idx])
+            i += idx
+        if len(labels_list) != len(pred_decode_labels):
+            raise Exception("origin labels length is {}, but pred labels length is {}".format(
+                len(labels_list), len(pred_decode_labels)))
+        for ids in range(len(labels_list)):
+            if labels_list[ids][0] == pred_decode_labels[ids].item():
+                correct_list.append(ids)
+            else:
+                error_list.append(ids)
+        return pred_decode_labels, labels_list, correct_list, error_list
+
+    def get_loss(self, predict, labels, labels_length):
+        labels = torch.autograd.Variable(labels)
+        if self.word:
+            loss = self.loss(predict, labels.long().cuda())
+        else:
+            log_predict = predict.log_softmax(2).detach().requires_grad_()
+            seq_len = torch.IntTensor([log_predict.shape[0]] * log_predict.shape[1])
+            loss = self.loss(log_predict.cpu(), labels, seq_len, labels_length)
+        self.optimizer.zero_grad()
+        loss.backward()
+        self.optimizer.step()
+
+        return loss.item(), self.scheduler.state_dict()['_last_lr'][-1]
+
+    def save_model(self, path, net):
+        torch.save(net, path)
+
+    def get_device(self, gpu_id):
+        if gpu_id == -1:
+            device = torch.device('cpu'.format(str(gpu_id)))
+        else:
+            device = torch.device('cuda:{}'.format(str(gpu_id)))
+        return device
+
+    def variable_to_device(self, inputs, device):
+        return torch.autograd.Variable(inputs).to(device)
+
+    def get_random_tensor(self):
+        width = self.resize[0]
+        height = self.resize[1]
+        if width == -1:
+            w = 240
+            h = height
+        else:
+            w = width
+            h = height
+        return torch.randn(1, self.image_channel, h, w, device='cpu')
+
+    def export_onnx(self, net, dummy_input, graph_path, input_names, output_names, dynamic_ax):
+        torch.onnx.export(net, dummy_input, graph_path, export_params=True, verbose=False,
+                          input_names=input_names, output_names=output_names, dynamic_axes=dynamic_ax,
+                          opset_version=12, do_constant_folding=True, _retain_param_name=False)

nets/backbone/__init__.py

@@ -0,0 +1,3 @@
+from .ddddocr import *
+from .effcientnet import *
+from .mobilenet import *

nets/backbone/ddddocr/__init__.py

@@ -0,0 +1 @@
+from .ddddocrv1 import DdddOcr

nets/backbone/ddddocr/ddddocrv1.py

@@ -0,0 +1,63 @@
+'''
+不记得从哪套模型改的了，可能来自于部门mobildenetv2
+'''
+import torch
+import torch.nn as nn
+
+
+class DdddOcr(nn.Module):
+    def __init__(self, nc=3, leakyRelu=False):
+        super(DdddOcr, self).__init__()
+        # assert imgH % 16 == 0, 'imgH has to be a multiple of 16'
+
+        ks = [3, 3, 3, 3, 3, 3, 2]
+        ps = [1, 1, 1, 1, 1, 1, 0]
+        ss = [1, 1, 1, 1, 1, 1, 1]
+        nm = [16, 32, 64, 64, 128, 128, 128]
+
+        cnn = nn.Sequential()
+
+        def convRelu(i, batchNormalization=False):
+            nIn = nc if i == 0 else nm[i - 1]
+            nOut = nm[i]
+            cnn.add_module('conv{0}'.format(i),
+                           nn.Conv2d(nIn, nOut, ks[i], ss[i], ps[i]))
+            if batchNormalization:
+                cnn.add_module('batchnorm{0}'.format(i), nn.BatchNorm2d(nOut))
+            if leakyRelu:
+                cnn.add_module('relu{0}'.format(i),
+                               nn.LeakyReLU(0.2, inplace=True))
+            else:
+                cnn.add_module('relu{0}'.format(i), nn.ReLU(True))
+
+        convRelu(0)
+        cnn.add_module('pooling{0}'.format(0), nn.MaxPool2d(2, 2))  # 64x16x64
+        convRelu(1)
+        cnn.add_module('pooling{0}'.format(1), nn.MaxPool2d(2, 2))  # 128x8x32
+        convRelu(2, True)
+        convRelu(3)
+        cnn.add_module('pooling{0}'.format(2),
+                       nn.MaxPool2d((2, 2), (2, 1), (0, 1)))  # 256x4x16
+        convRelu(4, True)
+        convRelu(5)
+        cnn.add_module('pooling{0}'.format(3),
+                       nn.MaxPool2d((2, 2), (2, 1), (0, 1)))  # 512x2x16
+        convRelu(6, True)  # 512x1x16
+
+        self.cnn = cnn
+
+    def forward(self, input):
+        return self.cnn(input)
+
+def test():
+    net = DdddOcr(1)
+    x = torch.randn(1, 1, 64, 224)
+    y = net(x)
+    print(y.size())
+    y = y.permute(3, 0, 1, 2)
+    w, b, c, h = y.shape
+    y = y.view(w, b, c * h)
+    print(y.size())
+
+if __name__ == '__main__':
+    test()

nets/backbone/effcientnet/__init__.py

@@ -0,0 +1 @@
+from .efficientnetv2 import effnetv2_l, effnetv2_m, effnetv2_xl, effnetv2_s

nets/backbone/effcientnet/efficientnetv2.py

@@ -0,0 +1,235 @@
+"""
+Creates a EfficientNetV2 Model as defined in:
+Mingxing Tan, Quoc V. Le. (2021).
+EfficientNetV2: Smaller Models and Faster Training
+arXiv preprint arXiv:2104.00298.
+import from https://github.com/d-li14/mobilenetv2.pytorch
+"""
+
+import torch
+import torch.nn as nn
+import math
+
+__all__ = ['effnetv2_s', 'effnetv2_m', 'effnetv2_l', 'effnetv2_xl']
+
+
+def _make_divisible(v, divisor, min_value=None):
+    """
+    This function is taken from the original tf repo.
+    It ensures that all layers have a channel number that is divisible by 8
+    It can be seen here:
+    https://github.com/tensorflow/models/blob/master/research/slim/nets/mobilenet/mobilenet.py
+    :param v:
+    :param divisor:
+    :param min_value:
+    :return:
+    """
+    if min_value is None:
+        min_value = divisor
+    new_v = max(min_value, int(v + divisor / 2) // divisor * divisor)
+    # Make sure that round down does not go down by more than 10%.
+    if new_v < 0.9 * v:
+        new_v += divisor
+    return new_v
+
+
+# SiLU (Swish) activation function
+if hasattr(nn, 'SiLU'):
+    SiLU = nn.SiLU
+else:
+    # For compatibility with old PyTorch versions
+    class SiLU(nn.Module):
+        def forward(self, x):
+            return x * torch.sigmoid(x)
+
+
+class SELayer(nn.Module):
+    def __init__(self, inp, oup, reduction=4):
+        super(SELayer, self).__init__()
+        self.avg_pool = nn.AdaptiveAvgPool2d(1)
+        self.fc = nn.Sequential(
+            nn.Linear(oup, _make_divisible(inp // reduction, 8)),
+            SiLU(),
+            nn.Linear(_make_divisible(inp // reduction, 8), oup),
+            nn.Sigmoid()
+        )
+
+    def forward(self, x):
+        b, c, _, _ = x.size()
+        y = self.avg_pool(x).view(b, c)
+        y = self.fc(y).view(b, c, 1, 1)
+        return x * y
+
+
+def conv_3x3_bn(inp, oup, stride):
+    return nn.Sequential(
+        nn.Conv2d(inp, oup, 3, stride, 1, bias=False),
+        nn.BatchNorm2d(oup),
+        SiLU()
+    )
+
+
+def conv_1x1_bn(inp, oup):
+    return nn.Sequential(
+        nn.Conv2d(inp, oup, 1, 1, 0, bias=False),
+        nn.BatchNorm2d(oup),
+        SiLU()
+    )
+
+
+class MBConv(nn.Module):
+    def __init__(self, inp, oup, stride, expand_ratio, use_se):
+        super(MBConv, self).__init__()
+        assert stride in [1, 2]
+
+        hidden_dim = round(inp * expand_ratio)
+        self.identity = stride == 1 and inp == oup
+        if use_se:
+            self.conv = nn.Sequential(
+                # pw
+                nn.Conv2d(inp, hidden_dim, 1, 1, 0, bias=False),
+                nn.BatchNorm2d(hidden_dim),
+                SiLU(),
+                # dw
+                nn.Conv2d(hidden_dim, hidden_dim, 3, stride, 1, groups=hidden_dim, bias=False),
+                nn.BatchNorm2d(hidden_dim),
+                SiLU(),
+                SELayer(inp, hidden_dim),
+                # pw-linear
+                nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),
+                nn.BatchNorm2d(oup),
+            )
+        else:
+            self.conv = nn.Sequential(
+                # fused
+                nn.Conv2d(inp, hidden_dim, 3, stride, 1, bias=False),
+                nn.BatchNorm2d(hidden_dim),
+                SiLU(),
+                # pw-linear
+                nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),
+                nn.BatchNorm2d(oup),
+            )
+
+    def forward(self, x):
+        if self.identity:
+            return x + self.conv(x)
+        else:
+            return self.conv(x)
+
+
+class EffNetV2(nn.Module):
+    def __init__(self, cfgs, nc=3, width_mult=1.):
+        super(EffNetV2, self).__init__()
+        self.cfgs = cfgs
+
+        # building first layer
+        input_channel = _make_divisible(24 * width_mult, 8)
+        layers = [conv_3x3_bn(nc, input_channel, 2)]
+        # building inverted residual blocks
+        block = MBConv
+        for t, c, n, s, use_se in self.cfgs:
+            output_channel = _make_divisible(c * width_mult, 8)
+            for i in range(n):
+                layers.append(block(input_channel, output_channel, s if i == 0 else 1, t, use_se))
+                input_channel = output_channel
+        self.features = nn.Sequential(*layers)
+
+        self._initialize_weights()
+
+    def forward(self, x):
+        x = self.features(x)
+        return x
+
+    def _initialize_weights(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
+                m.weight.data.normal_(0, math.sqrt(2. / n))
+                if m.bias is not None:
+                    m.bias.data.zero_()
+            elif isinstance(m, nn.BatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+            elif isinstance(m, nn.Linear):
+                m.weight.data.normal_(0, 0.001)
+                m.bias.data.zero_()
+
+
+def effnetv2_s(**kwargs):
+    """
+    Constructs a EfficientNetV2-S model
+    """
+    cfgs = [
+        # t, c, n, s, SE
+        [1, 24, 2, 1, 0],
+        [4, 48, 4, 2, 0],
+        [4, 64, 4, 2, 0],
+        [4, 128, 6, 2, 1],
+        [6, 160, 9, 1, 1],
+        [6, 256, 15, 2, 1],
+    ]
+    return EffNetV2(cfgs, **kwargs)
+
+
+def effnetv2_m(**kwargs):
+    """
+    Constructs a EfficientNetV2-M model
+    """
+    cfgs = [
+        # t, c, n, s, SE
+        [1, 24, 3, 1, 0],
+        [4, 48, 5, 2, 0],
+        [4, 80, 5, 2, 0],
+        [4, 160, 7, 2, 1],
+        [6, 176, 14, 1, 1],
+        [6, 304, 18, 2, 1],
+        [6, 512, 5, 1, 1],
+    ]
+    return EffNetV2(cfgs, **kwargs)
+
+
+def effnetv2_l(**kwargs):
+    """
+    Constructs a EfficientNetV2-L model
+    """
+    cfgs = [
+        # t, c, n, s, SE
+        [1, 32, 4, 1, 0],
+        [4, 64, 7, 2, 0],
+        [4, 96, 7, 2, 0],
+        [4, 192, 10, 2, 1],
+        [6, 224, 19, 1, 1],
+        [6, 384, 25, 2, 1],
+        [6, 640, 7, 1, 1],
+    ]
+    return EffNetV2(cfgs, **kwargs)
+
+
+def effnetv2_xl(**kwargs):
+    """
+    Constructs a EfficientNetV2-XL model
+    """
+    cfgs = [
+        # t, c, n, s, SE
+        [1, 32, 4, 1, 0],
+        [4, 64, 8, 2, 0],
+        [4, 96, 8, 2, 0],
+        [4, 192, 16, 2, 1],
+        [6, 256, 24, 1, 1],
+        [6, 512, 32, 2, 1],
+        [6, 640, 8, 1, 1],
+    ]
+    return EffNetV2(cfgs, **kwargs)
+
+def test():
+    net = effnetv2_s(nc=1)
+    x = torch.randn(2, 3, 50, 224)
+    y = net(x)
+    print(y.size())
+    y = y.permute(3, 0, 1, 2)
+    w, b, c, h = y.shape
+    y = y.view(w, b, c * h)
+    print(y.size())
+
+if __name__ == '__main__':
+    test()

nets/backbone/mobilenet/__init__.py

@@ -0,0 +1,2 @@
+from .mobilenetv2 import mobilenetv2
+from .mobilenetv3 import MobileNetV3_Small, MobileNetV3_Large

nets/backbone/mobilenet/mobilenetv2.py

@@ -0,0 +1,142 @@
+"""
+Creates a MobileNetV2 Model as defined in:
+Mark Sandler, Andrew Howard, Menglong Zhu, Andrey Zhmoginov, Liang-Chieh Chen. (2018).
+MobileNetV2: Inverted Residuals and Linear Bottlenecks
+arXiv preprint arXiv:1801.04381.
+import from https://github.com/tonylins/pytorch-mobilenet-v2
+"""
+
+import torch.nn as nn
+import math
+
+__all__ = ['mobilenetv2']
+
+
+def _make_divisible(v, divisor, min_value=None):
+    """
+    This function is taken from the original tf repo.
+    It ensures that all layers have a channel number that is divisible by 8
+    It can be seen here:
+    https://github.com/tensorflow/models/blob/master/research/slim/nets/mobilenet/mobilenet.py
+    :param v:
+    :param divisor:
+    :param min_value:
+    :return:
+    """
+    if min_value is None:
+        min_value = divisor
+    new_v = max(min_value, int(v + divisor / 2) // divisor * divisor)
+    # Make sure that round down does not go down by more than 10%.
+    if new_v < 0.9 * v:
+        new_v += divisor
+    return new_v
+
+
+def conv_3x3_bn(inp, oup, stride):
+    return nn.Sequential(
+        nn.Conv2d(inp, oup, 3, stride, 1, bias=False),
+        nn.BatchNorm2d(oup),
+        nn.ReLU6(inplace=True)
+    )
+
+
+def conv_1x1_bn(inp, oup):
+    return nn.Sequential(
+        nn.Conv2d(inp, oup, 1, 1, 0, bias=False),
+        nn.BatchNorm2d(oup),
+        nn.ReLU6(inplace=True)
+    )
+
+
+class InvertedResidual(nn.Module):
+    def __init__(self, inp, oup, stride, expand_ratio):
+        super(InvertedResidual, self).__init__()
+        assert stride in [1, 2]
+
+        hidden_dim = round(inp * expand_ratio)
+        self.identity = stride == 1 and inp == oup
+
+        if expand_ratio == 1:
+            self.conv = nn.Sequential(
+                # dw
+                nn.Conv2d(hidden_dim, hidden_dim, 3, stride, 1, groups=hidden_dim, bias=False),
+                nn.BatchNorm2d(hidden_dim),
+                nn.ReLU6(inplace=True),
+                # pw-linear
+                nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),
+                nn.BatchNorm2d(oup),
+            )
+        else:
+            self.conv = nn.Sequential(
+                # pw
+                nn.Conv2d(inp, hidden_dim, 1, 1, 0, bias=False),
+                nn.BatchNorm2d(hidden_dim),
+                nn.ReLU6(inplace=True),
+                # dw
+                nn.Conv2d(hidden_dim, hidden_dim, 3, stride, 1, groups=hidden_dim, bias=False),
+                nn.BatchNorm2d(hidden_dim),
+                nn.ReLU6(inplace=True),
+                # pw-linear
+                nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),
+                nn.BatchNorm2d(oup),
+            )
+
+    def forward(self, x):
+        if self.identity:
+            return x + self.conv(x)
+        else:
+            return self.conv(x)
+
+
+class MobileNetV2(nn.Module):
+    def __init__(self, nc=3, width_mult=1.):
+        super(MobileNetV2, self).__init__()
+        # setting of inverted residual blocks
+        self.cfgs = [
+            # t, c, n, s
+            [1,  16, 1, 1],
+            [6,  24, 2, 2],
+            [6,  32, 3, 2],
+            [6,  64, 4, 2],
+            [6,  96, 3, 1],
+            [6, 160, 3, 2],
+            [6, 320, 1, 1],
+        ]
+
+        # building first layer
+        input_channel = _make_divisible(32 * width_mult, 4 if width_mult == 0.1 else 8)
+        layers = [conv_3x3_bn(nc, input_channel, 2)]
+        # building inverted residual blocks
+        block = InvertedResidual
+        for t, c, n, s in self.cfgs:
+            output_channel = _make_divisible(c * width_mult, 4 if width_mult == 0.1 else 8)
+            for i in range(n):
+                layers.append(block(input_channel, output_channel, s if i == 0 else 1, t))
+                input_channel = output_channel
+        self.features = nn.Sequential(*layers)
+
+        self._initialize_weights()
+
+    def forward(self, x):
+        x = self.features(x)
+        return x
+
+    def _initialize_weights(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
+                m.weight.data.normal_(0, math.sqrt(2. / n))
+                if m.bias is not None:
+                    m.bias.data.zero_()
+            elif isinstance(m, nn.BatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+            elif isinstance(m, nn.Linear):
+                m.weight.data.normal_(0, 0.01)
+                m.bias.data.zero_()
+
+def mobilenetv2(**kwargs):
+    """
+    Constructs a MobileNet V2 model
+    """
+    return MobileNetV2(**kwargs)

nets/backbone/mobilenet/mobilenetv3.py

@@ -0,0 +1,191 @@
+'''MobileNetV3 in PyTorch.
+See the paper "Inverted Residuals and Linear Bottlenecks:
+Mobile Networks for Classification, Detection and Segmentation" for more details.
+
+import from https://github.com/xiaolai-sqlai/mobilenetv3/blob/master/mobilenetv3.py
+'''
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn import init
+
+
+class hswish(nn.Module):
+    def forward(self, x):
+        out = x * F.relu6(x + 3, inplace=True) / 6
+        return out
+
+
+class hsigmoid(nn.Module):
+    def forward(self, x):
+        out = F.relu6(x + 3, inplace=True) / 6
+        return out
+
+
+class SeModule(nn.Module):
+    def __init__(self, in_size, reduction=4):
+        super(SeModule, self).__init__()
+        self.se = nn.Sequential(
+            nn.AdaptiveAvgPool2d(1),
+            nn.Conv2d(in_size, in_size // reduction, kernel_size=1, stride=1, padding=0, bias=False),
+            nn.BatchNorm2d(in_size // reduction),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(in_size // reduction, in_size, kernel_size=1, stride=1, padding=0, bias=False),
+            nn.BatchNorm2d(in_size),
+            hsigmoid()
+        )
+
+    def forward(self, x):
+        return x * self.se(x)
+
+
+class Block(nn.Module):
+    '''expand + depthwise + pointwise'''
+
+    def __init__(self, kernel_size, in_size, expand_size, out_size, nolinear, semodule, stride):
+        super(Block, self).__init__()
+        self.stride = stride
+        self.se = semodule
+
+        self.conv1 = nn.Conv2d(in_size, expand_size, kernel_size=1, stride=1, padding=0, bias=False)
+        self.bn1 = nn.BatchNorm2d(expand_size)
+        self.nolinear1 = nolinear
+        self.conv2 = nn.Conv2d(expand_size, expand_size, kernel_size=kernel_size, stride=stride,
+                               padding=kernel_size // 2, groups=expand_size, bias=False)
+        self.bn2 = nn.BatchNorm2d(expand_size)
+        self.nolinear2 = nolinear
+        self.conv3 = nn.Conv2d(expand_size, out_size, kernel_size=1, stride=1, padding=0, bias=False)
+        self.bn3 = nn.BatchNorm2d(out_size)
+
+        self.shortcut = nn.Sequential()
+        if stride == 1 and in_size != out_size:
+            self.shortcut = nn.Sequential(
+                nn.Conv2d(in_size, out_size, kernel_size=1, stride=1, padding=0, bias=False),
+                nn.BatchNorm2d(out_size),
+            )
+
+    def forward(self, x):
+        out = self.nolinear1(self.bn1(self.conv1(x)))
+        out = self.nolinear2(self.bn2(self.conv2(out)))
+        out = self.bn3(self.conv3(out))
+        if self.se != None:
+            out = self.se(out)
+        out = out + self.shortcut(x) if self.stride == 1 else out
+        return out
+
+
+class MobileNetV3_Large(nn.Module):
+    def __init__(self, nc=3):
+        super(MobileNetV3_Large, self).__init__()
+        self.conv1 = nn.Conv2d(nc, 16, kernel_size=3, stride=2, padding=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(16)
+        self.hs1 = hswish()
+
+        self.bneck = nn.Sequential(
+            Block(3, 16, 16, 16, nn.ReLU(inplace=True), None, 1),
+            Block(3, 16, 64, 24, nn.ReLU(inplace=True), None, 2),
+            Block(3, 24, 72, 24, nn.ReLU(inplace=True), None, 1),
+            Block(5, 24, 72, 40, nn.ReLU(inplace=True), SeModule(40), 2),
+            Block(5, 40, 120, 40, nn.ReLU(inplace=True), SeModule(40), 1),
+            Block(5, 40, 120, 40, nn.ReLU(inplace=True), SeModule(40), 1),
+            Block(3, 40, 240, 80, hswish(), None, 2),
+            Block(3, 80, 200, 80, hswish(), None, 1),
+            Block(3, 80, 184, 80, hswish(), None, 1),
+            Block(3, 80, 184, 80, hswish(), None, 1),
+            Block(3, 80, 480, 112, hswish(), SeModule(112), 1),
+            Block(3, 112, 672, 112, hswish(), SeModule(112), 1),
+            Block(5, 112, 672, 160, hswish(), SeModule(160), 1),
+            Block(5, 160, 672, 160, hswish(), SeModule(160), 2),
+            Block(5, 160, 960, 160, hswish(), SeModule(160), 1),
+        )
+
+        self.conv2 = nn.Conv2d(160, 960, kernel_size=1, stride=1, padding=0, bias=False)
+        self.bn2 = nn.BatchNorm2d(960)
+        self.hs2 = hswish()
+        self.linear3 = nn.Linear(960, 1280)
+        self.bn3 = nn.BatchNorm1d(1280)
+        self.init_params()
+
+    def init_params(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                init.kaiming_normal_(m.weight, mode='fan_out')
+                if m.bias is not None:
+                    init.constant_(m.bias, 0)
+            elif isinstance(m, nn.BatchNorm2d):
+                init.constant_(m.weight, 1)
+                init.constant_(m.bias, 0)
+            elif isinstance(m, nn.Linear):
+                init.normal_(m.weight, std=0.001)
+                if m.bias is not None:
+                    init.constant_(m.bias, 0)
+
+    def forward(self, x):
+        out = self.hs1(self.bn1(self.conv1(x)))
+        out = self.bneck(out)
+        out = self.hs2(self.bn2(self.conv2(out)))
+        return out
+
+
+class MobileNetV3_Small(nn.Module):
+    def __init__(self, nc=3):
+        super(MobileNetV3_Small, self).__init__()
+        self.conv1 = nn.Conv2d(nc, 16, kernel_size=3, stride=2, padding=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(16)
+        self.hs1 = hswish()
+
+        self.bneck = nn.Sequential(
+            Block(3, 16, 16, 16, nn.ReLU(inplace=True), SeModule(16), 2),
+            Block(3, 16, 72, 24, nn.ReLU(inplace=True), None, 2),
+            Block(3, 24, 88, 24, nn.ReLU(inplace=True), None, 1),
+            Block(5, 24, 96, 40, hswish(), SeModule(40), 2),
+            Block(5, 40, 240, 40, hswish(), SeModule(40), 1),
+            Block(5, 40, 240, 40, hswish(), SeModule(40), 1),
+            Block(5, 40, 120, 48, hswish(), SeModule(48), 1),
+            Block(5, 48, 144, 48, hswish(), SeModule(48), 1),
+            Block(5, 48, 288, 96, hswish(), SeModule(96), 2),
+            Block(5, 96, 576, 96, hswish(), SeModule(96), 1),
+            Block(5, 96, 576, 96, hswish(), SeModule(96), 1),
+        )
+
+        self.conv2 = nn.Conv2d(96, 576, kernel_size=1, stride=1, padding=0, bias=False)
+        self.bn2 = nn.BatchNorm2d(576)
+        self.hs2 = hswish()
+        self.linear3 = nn.Linear(576, 1280)
+        self.bn3 = nn.BatchNorm1d(1280)
+        self.init_params()
+
+    def init_params(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                init.kaiming_normal_(m.weight, mode='fan_out')
+                if m.bias is not None:
+                    init.constant_(m.bias, 0)
+            elif isinstance(m, nn.BatchNorm2d):
+                init.constant_(m.weight, 1)
+                init.constant_(m.bias, 0)
+            elif isinstance(m, nn.Linear):
+                init.normal_(m.weight, std=0.001)
+                if m.bias is not None:
+                    init.constant_(m.bias, 0)
+
+    def forward(self, x):
+        out = self.hs1(self.bn1(self.conv1(x)))
+        out = self.bneck(out)
+        out = self.hs2(self.bn2(self.conv2(out)))
+
+        return out
+
+
+def test():
+    net = MobileNetV3_Small()
+    x = torch.randn(2, 3, 50, 224)
+    y = net(x)
+    print(y.size())
+    y = y.permute(3, 0, 1, 2)
+    w, b, c, h = y.shape
+    y = y.view(w, b, c * h)
+    print(y.size())
+
+if __name__ == '__main__':
+    test()

requirements.txt

@@ -0,0 +1,9 @@
+fire~=0.4.0
+loguru~=0.5.3
+yaml~=0.2.5
+pyyaml~=6.0
+torch~=1.10.0+cu113
+tqdm~=4.62.3
+numpy~=1.20.3
+torchvision~=0.11.1+cu113
+ddddocr~=1.3.0

utils/cache_data.py

@@ -0,0 +1,116 @@
+import json
+import os
+import random
+
+import tqdm
+
+from configs import Config
+from loguru import logger
+
+
+class CacheData:
+    def __init__(self, project_name: str):
+        self.project_name = project_name
+        self.project_path = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "projects",
+                                         project_name)
+        if os.path.exists(self.project_path):
+            self.cache_path = os.path.join(self.project_path, "cache")
+        else:
+            logger.error("Project {} is not exists!".format(project_name))
+            exit()
+        self.config = Config(project_name)
+        self.conf = self.config.load_config()
+        self.bath_path = self.conf['System']['Path']
+        self.allow_ext = []
+
+    def cache(self, base_path: str, search_type="name"):
+        self.bath_path = base_path
+        self.allow_ext = self.conf["System"]["Allow_Ext"]
+        if search_type == "name":
+            self.__get_label_from_name(base_path=base_path)
+        else:
+            self.__get_label_from_file(base_path=base_path)
+
+    def __get_label_from_name(self, base_path: str):
+        files = os.listdir(base_path)
+        logger.info("\nFiles number is {}.".format(len(files)))
+        self.__collect_data(files, base_path)
+
+    def __get_label_from_file(self, base_path: str):
+        labels_path = os.path.join(base_path, "labels.txt")
+        images_path = os.path.join(base_path, "images")
+        if not os.path.exists(labels_path):
+            logger.error("\nThe file labels.txt not found in path ----> {}".format(base_path))
+            exit()
+        if not os.path.exists(images_path) or not os.path.isdir(images_path):
+            logger.error("\nThe dir {} not found in path ----> {}".format(images_path, base_path))
+            exit()
+        files = os.listdir(images_path)
+        logger.info("\nFiles number is {}.".format(len(files)))
+        with open(labels_path, "r", encoding="utf-8") as f:
+            labels_lines = f.readlines()
+        labels_lines = [line.replace("\r", "").replace("\n", "") for line in labels_lines]
+        labels_filename_lines = [line.split("\t")[0] for line in labels_lines]
+        logger.info("\nLabels number is {}.".format(len(labels_lines)))
+        logger.info("\nChecking labels.txt ...")
+        error_files = set(labels_filename_lines).difference(set(files))
+        logger.info("\nCheck labels.txt end! {} errors!".format(len(error_files)))
+        for ef in error_files:
+            labels_lines.remove(ef)
+        del files
+        self.__collect_data(labels_lines, images_path, is_file=True)
+
+    def __collect_data(self, lines, base_path, is_file=False):
+        labels = []
+        caches = []
+        if not self.conf['Model']['Word']:
+            labels.append(" ")
+        for file in tqdm.tqdm(lines):
+            if is_file:
+                line_list = file.split('\t')
+                filename = line_list[0]
+                label = line_list[1]
+            else:
+                filename = file
+                label = "_".join(filename.split("_")[:-1])
+            label = label.replace(" ", "")
+            if filename.split('.')[-1] in self.allow_ext:
+                if " " in filename:
+                    logger.warning("The {} has black. We will remove it!".format(filename))
+                    continue
+                caches.append('\t'.join([filename, label]))
+                if not self.conf['Model']['Word']:
+                    label = list(label)
+                    labels.extend(label)
+                else:
+                    labels.append(label)
+
+            else:
+                logger.warning("\nFile({}) has a suffix that is not allowed! We will remove it!".format(file))
+        labels = list(set(labels))
+        logger.info("\nCoolect labels is {}".format(json.dumps(labels, ensure_ascii=False)))
+        self.conf['System']['Path'] = base_path
+        self.conf['Model']['CharSet'] = labels
+        self.config.make_config(config_dict=self.conf, single=self.conf['Model']['Word'])
+        logger.info("\nWriting Cache Data!")
+        del lines
+        logger.info("\nCache Data Number is {}".format(len(caches)))
+        logger.info("\nWriting Train and Val File.".format(len(caches)))
+        val = self.conf['System']['Val']
+        if 0 < val < 1:
+            val_num = int(len(caches) * val)
+        elif 1 < val < len(caches):
+            val_num = int(val)
+        else:
+            logger.error("val setting vaild!")
+            exit()
+        random.shuffle(caches)
+        train_set = caches[val_num:]
+        val_set = caches[:val_num]
+        del caches
+        with open(os.path.join(self.cache_path, "cache.train.tmp"), 'w', encoding="utf-8") as f:
+            f.write("\n".join(train_set))
+        with open(os.path.join(self.cache_path, "cache.val.tmp"), 'w', encoding="utf-8") as f:
+            f.write("\n".join(val_set))
+        logger.info("\nTrain Data Number is {}".format(len(train_set)))
+        logger.info("\nVal Data Number is {}".format(len(val_set)))

utils/load_cache.py

@@ -0,0 +1,162 @@
+import json
+import os
+
+import torch
+import tqdm
+import numpy as np
+
+from configs import Config
+from loguru import logger
+
+import torchvision
+from torch.utils.data import DataLoader, Dataset, TensorDataset
+
+
+class LoadCache(Dataset):
+    def __init__(self, cache_path: str, path: str, word: bool, image_channel: int, resize: list, charset: list):
+        self.cache_path = cache_path
+        self.path = path
+        self.word = word
+        self.ImageChannel = image_channel
+        self.resize = resize
+        self.charset = charset
+
+        logger.info("\nReading Cache File... ----> {}".format(self.cache_path))
+
+        with open(self.cache_path, 'r', encoding='utf-8') as f:
+            caches = f.readlines()
+        self.caches = []
+        for cache in tqdm.tqdm(caches):
+            cache = cache.replace("\r", "").replace("\n", "").split("\t")
+            self.caches.append(cache)
+        del caches
+
+        self.caches_num = len(self.caches)
+        logger.info("\nRead Cache File End! Caches Num is {}.".format(self.caches_num))
+
+    def __len__(self):
+        return self.caches_num
+
+    def __getitem__(self, idx):
+        try:
+            data = self.caches[idx]
+            image_name = data[0]
+            image_label = data[1]
+            image_path = os.path.join(self.path, image_name)
+            if not self.word:
+                image_label = list(image_label)
+            else:
+                image_label = [image_label]
+            if self.ImageChannel == 1:
+                mode = torchvision.io.ImageReadMode.GRAY
+            else:
+                mode = torchvision.io.ImageReadMode.RGB
+            image = torchvision.io.read_image(image_path, mode=mode)  # shape c, h, w
+            image_shape = image.shape
+            image_height = image_shape[1]
+            image_width = image_shape[2]
+            width = self.resize[0]
+            height = self.resize[1]
+            if self.resize[0] == -1:
+                image = torchvision.transforms.Resize((height, int(image_width * (height / image_height))))(image)
+            else:
+                image = torchvision.transforms.Resize((height, width))(image)
+            image = torchvision.transforms.ToPILImage()(image)
+            label = [int(self.charset.index(item)) for item in list(image_label)]
+            return image, label
+
+        except Exception as e:
+            logger.error("\nError: {}, File: {}".format(str(e), self.caches[idx][0]))
+            return None, None
+
+
+class GetLoader:
+    def __init__(self, project_name: str):
+        self.project_name = project_name
+        self.project_path = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "projects",
+                                         project_name)
+        if os.path.exists(self.project_path):
+            self.cache_path = os.path.join(self.project_path, "cache")
+            if os.path.exists(self.cache_path):
+                self.cache_train_path = os.path.join(self.cache_path, "cache.train.tmp")
+                self.cache_val_path = os.path.join(self.cache_path, "cache.val.tmp")
+
+                if not os.path.exists(self.cache_train_path):
+                    logger.error("\nCache Train File {} is not exists!".format(self.cache_train_path))
+                    exit()
+                if not os.path.exists(self.cache_val_path):
+                    logger.error("\nCache Val File {} is not exists!".format(self.cache_val_path))
+                    exit()
+
+            else:
+                logger.error("\nCache dir {} is not exists!".format(self.cache_path))
+                exit()
+        else:
+            logger.error("\nProject {} is not exists!".format(project_name))
+            exit()
+
+        self.config = Config(project_name)
+        self.conf = self.config.load_config()
+
+        self.charset = self.conf['Model']['CharSet']
+        logger.info("\nCharsets is {}".format(json.dumps(self.charset, ensure_ascii=False)))
+
+        self.resize = [int(self.conf['Model']['ImageWidth']), int(self.conf['Model']['ImageHeight'])]
+        logger.info("\nImage Resize is {}".format(json.dumps(self.resize)))
+
+        self.ImageChannel = self.conf['Model']['ImageChannel']
+
+        self.word = self.conf['Model']['Word']
+
+        self.path = self.conf['System']['Path']
+
+        self.batch_size = self.conf['Train']['BATCH_SIZE']
+
+        self.val_batch_size = self.conf['Train']['TEST_BATCH_SIZE']
+
+        logger.info("\nImage Path is {}".format(self.path))
+
+        self.transform_list = []
+        self.transform_list.append(torchvision.transforms.ToTensor())
+        if self.ImageChannel == 1:
+            self.transform_list.append(torchvision.transforms.Normalize(mean=[0.456],
+                                                                   std=[0.224]))
+        else:
+            if self.ImageChannel != 3:
+                logger.error("ImageChannel must be 1 or 3!")
+                exit()
+            self.transform_list.append(torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                                                                   std=[0.229, 0.224, 0.225]))
+        self.transform = torchvision.transforms.Compose(self.transform_list)
+        tarin_loader = LoadCache(self.cache_train_path, self.path, self.word, self.ImageChannel, self.resize, self.charset)
+        val_loader = LoadCache(self.cache_val_path, self.path, self.word, self.ImageChannel, self.resize, self.charset)
+        self.loaders = {
+            'train': DataLoader(dataset=tarin_loader, batch_size=self.batch_size, shuffle=True, drop_last=True,
+                                num_workers=0, collate_fn=self.collate_to_sparse),
+            'val': DataLoader(dataset=val_loader, batch_size=self.val_batch_size, shuffle=True, drop_last=True,
+                                num_workers=0, collate_fn=self.collate_to_sparse),
+        }
+
+    def collate_to_sparse(self, batch):
+        values = []
+        images = []
+        shapes = []
+        max_width = 0
+        for n, (img, seq) in enumerate(batch):
+            if img is None or seq is None:
+                continue
+            if len(seq) == 0: continue
+            if max_width < img.size[0]:
+                max_width = img.size[0]
+            values.extend(seq)
+            images.append(img)
+            shapes.append(len(seq))
+        images_pad = []
+        for img in images:
+            img = torchvision.transforms.Pad((0, 0, int(max_width - img.size[0]), 0))(img)
+            if self.transform is not None:
+                img = self.transform(img)
+            images_pad.append(img)
+        images_pad = torch.stack(images_pad, dim=0)
+        return [images_pad, torch.FloatTensor(values), torch.IntTensor(shapes)]
+

utils/project_manager.py

@@ -0,0 +1,39 @@
+import os
+from configs import Config
+from loguru import logger
+
+
+class ProjectManager:
+
+    def __init__(self):
+        self.base_path = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "projects")
+
+    def create_project(self, project_name: str, single: bool = False):
+        project_base_path = os.path.join(self.base_path, project_name)
+        logger.info("Creating Directory... ----> {}".format(project_base_path))
+        if not os.path.exists(project_base_path):
+            os.mkdir(project_base_path)
+            if not os.path.exists(project_base_path):
+                logger.error("Directory create failed! ----> {}".format(project_base_path))
+                return False
+            models_path = os.path.join(project_base_path, "models")
+            logger.info("Creating Directory... ----> {}".format(models_path))
+            os.mkdir(models_path)
+
+            cache_path = os.path.join(project_base_path, "cache")
+            logger.info("Creating Directory... ----> {}".format(cache_path))
+            os.mkdir(cache_path)
+
+            checkpoints_path = os.path.join(project_base_path, "checkpoints")
+            logger.info("Creating Directory... ----> {}".format(checkpoints_path))
+            os.mkdir(checkpoints_path)
+
+            config_path = os.path.join(os.path.join(project_base_path, "config.yaml"))
+            logger.info("Creating {} Config File... ----> {}".format("CNN" if single else "CRNN", config_path))
+            conf = Config(project_name)
+            conf.make_config(single=single)
+
+            logger.info("Create Project Success! ----> {}".format(project_name))
+        else:
+            logger.error("Directory already exists! ----> {}".format(project_base_path))
+            return False

utils/train.py

@@ -0,0 +1,128 @@
+import json
+import os
+import random
+import time
+
+import tqdm
+
+from configs import Config
+from loguru import logger
+from utils import load_cache
+from nets import Net
+
+
+class Train:
+    def __init__(self, project_name: str):
+        self.project_name = project_name
+        self.project_path = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "projects",
+                                         project_name)
+        self.checkpoints_path = os.path.join(self.project_path, "checkpoints")
+        self.models_path = os.path.join(self.project_path, "models")
+        self.config = Config(project_name)
+        self.conf = self.config.load_config()
+
+        self.test_step = self.conf['Train']['TEST_STEP']
+        self.target = self.conf['Train']['TARGET']
+        self.target_acc = self.target['Accuracy']
+        self.min_epoch = self.target['Epoch']
+        self.max_loss = self.target['Cost']
+        logger.info("\nTaget:\nmin_Accuracy: {}\nmin_Epoch: {}\nmax_Loss: {}".format(self.target_acc, self.min_epoch,
+                                                                                     self.max_loss))
+
+        logger.info("\nBuilding Net...")
+        self.net = Net(self.conf)
+        logger.info(self.net)
+        logger.info("\nBuilding End")
+
+        self.use_gpu = self.conf['System']['GPU']
+        if self.use_gpu:
+            self.gpu_id = self.conf['System']['GPU_ID']
+            logger.info("\nUSE GPU ----> {}".format(self.gpu_id))
+            self.device = self.net.get_device(self.gpu_id)
+            self.net.to(self.device)
+        else:
+            self.gpu_id = -1
+            self.device = self.net.get_device(self.gpu_id)
+            logger.info("\nUSE CPU".format(self.gpu_id))
+        logger.info("\nGet Data Loader...")
+        loaders = load_cache.GetLoader(project_name)
+        self.train = loaders.loaders['train']
+        self.val = loaders.loaders['val']
+        logger.info("\nGet Data Loader End!")
+
+        self.epoch = 0
+        self.step = 0
+        self.loss = 0
+        self.avg_loss = 0
+        self.start_time = time.time()
+        self.now_time = time.time()
+
+    def start(self):
+        val_iter = iter(self.val)
+        while True:
+            for idx, (inputs, labels, labels_length) in enumerate(self.train):
+                self.now_time = time.time()
+                inputs = self.net.variable_to_device(inputs, device=self.device)
+
+                loss, lr = self.net.training(inputs, labels, labels_length)
+
+                self.avg_loss += loss
+
+                self.step += 1
+
+                if self.step % 100 == 0 and self.step % self.test_step != 0:
+                    logger.info("{}\tEpoch: {}\tStep: {}\tLastLoss: {}\tAvgLoss: {}\tLr: {}".format(
+                        time.strftime("[%Y-%m-%d-%H_%M_%S]", time.localtime(self.now_time)), self.epoch, self.step,
+                        str(loss), str(self.avg_loss / 100), lr
+                    ))
+                    self.avg_loss = 0
+                if self.step % 2000 == 0 and self.step != 0:
+                    model_path = os.path.join(self.checkpoints_path, "checkpoint_{}_{}_{}.tar".format(
+                        self.project_name, self.epoch, self.step,
+                    ))
+                    self.net.save_model(model_path,
+                                        {"net": self.net.state_dict(), "optimizer": self.net.optimizer.state_dict(),
+                                         "epoch": self.epoch, "step": self.step})
+                if self.step % self.test_step == 0:
+                    try:
+                        test_inputs, test_labels, test_labels_length = next(val_iter)
+                    except Exception:
+                        del val_iter
+                        val_iter = iter(self.val)
+                        test_inputs, test_labels, test_labels_length = next(val_iter)
+                    if test_inputs.shape[0] < 5:
+                        continue
+                    test_inputs = self.net.variable_to_device(test_inputs, self.device)
+                    self.net = self.net.train(False)
+                    pred_labels, labels_list, correct_list, error_list = self.net.test_op(test_inputs, test_labels,
+                                                                                          test_labels_length)
+                    self.net = self.net.train()
+                    accuracy = len(correct_list) / test_inputs.shape[0]
+                    logger.info("{}\tEpoch: {}\tStep: {}\tLastLoss: {}\tAvgLoss: {}\tLr: {}\tAcc: {}".format(
+                        time.strftime("[%Y-%m-%d-%H_%M_%S]", time.localtime(self.now_time)), self.epoch, self.step,
+                        str(loss), str(self.avg_loss / 100), lr, accuracy
+                    ))
+                    if accuracy > self.target_acc and self.epoch > self.min_epoch and self.avg_loss < self.max_loss:
+                        logger.info("\nTraining Finished!Exporting Model...")
+                        dummy_input = self.net.get_random_tensor()
+                        input_names = ["input1"]
+                        output_names = ["output"]
+
+                        if self.net.backbone.startswith("effnet"):
+                            self.net.cnn.set_swish(memory_efficient=False)
+                        self.net = self.net.eval().cpu()
+                        dynamic_ax = {'input1': {3: 'image_wdith'}, "output": {1: 'seq'}}
+                        self.net.export_onnx(self.net, dummy_input,
+                                             os.path.join(self.models_path, "{}_{}_{}_{}_{}.onnx".format(
+                                                 self.project_name, str(accuracy), self.epoch, self.step,
+                                                 time.localtime(self.now_time)))
+                                             , input_names, output_names, dynamic_ax)
+                        logger.info("\nExport Finished!Using Time: {}min".format(str(int(int(self.now_time * 1000) - int(self.start_time * 1000)) / 60)))
+                        exit()
+
+            self.epoch += 1
+            self.net.scheduler.step(self.epoch)
+
+
+if __name__ == '__main__':
+    Train("test1")