diff --git a/app.py b/app.py new file mode 100644 index 0000000..83a6d36 --- /dev/null +++ b/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) diff --git a/configs/__init__.py b/configs/__init__.py new file mode 100644 index 0000000..7556db6 --- /dev/null +++ b/configs/__init__.py @@ -0,0 +1,3 @@ +from .base import * + + diff --git a/configs/base.py b/configs/base.py new file mode 100644 index 0000000..4ef59ad --- /dev/null +++ b/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 diff --git a/nets/__init__.py b/nets/__init__.py new file mode 100644 index 0000000..1141d2e --- /dev/null +++ b/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) \ No newline at end of file diff --git a/nets/backbone/__init__.py b/nets/backbone/__init__.py new file mode 100644 index 0000000..0903739 --- /dev/null +++ b/nets/backbone/__init__.py @@ -0,0 +1,3 @@ +from .ddddocr import * +from .effcientnet import * +from .mobilenet import * \ No newline at end of file diff --git a/nets/backbone/ddddocr/__init__.py b/nets/backbone/ddddocr/__init__.py new file mode 100644 index 0000000..7c003d9 --- /dev/null +++ b/nets/backbone/ddddocr/__init__.py @@ -0,0 +1 @@ +from .ddddocrv1 import DdddOcr \ No newline at end of file diff --git a/nets/backbone/ddddocr/ddddocrv1.py b/nets/backbone/ddddocr/ddddocrv1.py new file mode 100644 index 0000000..b8c6468 --- /dev/null +++ b/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() \ No newline at end of file diff --git a/nets/backbone/effcientnet/__init__.py b/nets/backbone/effcientnet/__init__.py new file mode 100644 index 0000000..48a0699 --- /dev/null +++ b/nets/backbone/effcientnet/__init__.py @@ -0,0 +1 @@ +from .efficientnetv2 import effnetv2_l, effnetv2_m, effnetv2_xl, effnetv2_s \ No newline at end of file diff --git a/nets/backbone/effcientnet/efficientnetv2.py b/nets/backbone/effcientnet/efficientnetv2.py new file mode 100644 index 0000000..1e56648 --- /dev/null +++ b/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() \ No newline at end of file diff --git a/nets/backbone/mobilenet/__init__.py b/nets/backbone/mobilenet/__init__.py new file mode 100644 index 0000000..0b19a7d --- /dev/null +++ b/nets/backbone/mobilenet/__init__.py @@ -0,0 +1,2 @@ +from .mobilenetv2 import mobilenetv2 +from .mobilenetv3 import MobileNetV3_Small, MobileNetV3_Large \ No newline at end of file diff --git a/nets/backbone/mobilenet/mobilenetv2.py b/nets/backbone/mobilenet/mobilenetv2.py new file mode 100644 index 0000000..95f2507 --- /dev/null +++ b/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) diff --git a/nets/backbone/mobilenet/mobilenetv3.py b/nets/backbone/mobilenet/mobilenetv3.py new file mode 100644 index 0000000..7a78a9e --- /dev/null +++ b/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() \ No newline at end of file diff --git a/projects/__init__.py b/projects/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..6c6cebb --- /dev/null +++ b/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 \ No newline at end of file diff --git a/tools/__init__.py b/tools/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/utils/__init__.py b/utils/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/utils/cache_data.py b/utils/cache_data.py new file mode 100644 index 0000000..c436ea7 --- /dev/null +++ b/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))) diff --git a/utils/load_cache.py b/utils/load_cache.py new file mode 100644 index 0000000..e39e4e0 --- /dev/null +++ b/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)] + diff --git a/utils/project_manager.py b/utils/project_manager.py new file mode 100644 index 0000000..fa5eb96 --- /dev/null +++ b/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 diff --git a/utils/train.py b/utils/train.py new file mode 100644 index 0000000..86227ff --- /dev/null +++ b/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")