achieve the accuracy 80.xx

Author: chenjun2hao

dataset.py

@@ -141,7 +141,7 @@ def __init__(self, root, opt):
                     label_key = 'label-%09d'.encode() % index
                     label = txn.get(label_key).decode('utf-8')
 
-                    if len(label) >= self.opt.batch_max_length or len(label) == 0:
+                    if len(label) > self.opt.batch_max_length or len(label) == 0:
                         # print(f'The length of the label is longer than max_length: length \
                         #     {len(label)}, {label} in dataset {self.root}')
                         continue

model.py

@@ -90,7 +90,7 @@ def __init__(self, opt):
         elif opt.Prediction == 'Bert_pred':
             pass
         elif opt.Prediction == 'SRN':
-            self.Prediction = SRN_Decoder(n_position=opt.position_dim, n_class=opt.alphabet_size)
+            self.Prediction = SRN_Decoder(n_position=opt.position_dim, N_max_character=opt.batch_max_character + 1, n_class=opt.alphabet_size)
         else:
             raise Exception('Prediction is neither CTC or Attn')
 
@@ -103,7 +103,7 @@ def forward(self, input, text, is_train=True):
         """ Feature extraction stage """
         visual_feature = self.FeatureExtraction(input)
         # if self.stages['Feat'] == 'AsterRes' or self.stages['Feat'] == 'ResnetFpn':
-        if self.stages['Feat'] == 'AsterRes' or self.stages['Feat'] == 'ResnetFpn' or self.stages['Feat'] == 'ResNet':
+        if self.stages['Feat'] == 'AsterRes' or self.stages['Feat'] == 'ResnetFpn':
             b, c, h, w = visual_feature.shape
             visual_feature = visual_feature.permute(0, 1, 3, 2)
             visual_feature = visual_feature.contiguous().view(b, c, -1)

modules/feature_extraction.py

@@ -1,3 +1,5 @@
+# coding:utf-8
+# 2020-05-11
 import torch.nn as nn
 import torch.nn.functional as F
 

test.py

@@ -234,8 +234,8 @@ def test(opt):
     parser.add_argument('--eval_data', default='/home/deepblue/deepbluetwo/chenjun/1_OCR/data/data_lmdb_release/evaluation', help='path to evaluation dataset')
     parser.add_argument('--benchmark_all_eval', default=True, help='evaluate 10 benchmark evaluation datasets')
     parser.add_argument('--workers', type=int, help='number of data loading workers', default=4)
-    parser.add_argument('--batch_size', type=int, default=192, help='input batch size')
-    parser.add_argument('--saved_model', default='./saved_models/None-ResNet-SRN-SRN-Seed666/iter_210000.pth', help="path to saved_model to evaluation")
+    parser.add_argument('--batch_size', type=int, default=64, help='input batch size')
+    parser.add_argument('--saved_model', default='./saved_models/None-ResNet-SRN-SRN-Seed666/iter_65000.pth', help="path to saved_model to evaluation")
     """ Data processing """
     parser.add_argument('--batch_max_length', type=int, default=25, help='maximum-label-length')
     parser.add_argument('--imgH', type=int, default=32, help='the height of the input image')
@@ -258,6 +258,7 @@ def test(opt):
     parser.add_argument('--position_dim', type=int, default=26, help='the length sequence out from cnn encoder,resnet:65;resnetfpn:256')
 
     parser.add_argument('--SRN_PAD', type=int, default=36, help='the pad character for srn')
+    parser.add_argument('--batch_max_character', type=int, default=25, help='the max sequence length')
     opt = parser.parse_args()
 
     """ vocab / character number configuration """

train.py

@@ -126,7 +126,7 @@ def train(opt):
     print("Optimizer:")
     print(optimizer)
 
-    lrScheduler = lr_scheduler.MultiStepLR(optimizer, [2, 5, 10], gamma=0.1)                         # 减小学习速率
+    lrScheduler = lr_scheduler.MultiStepLR(optimizer, [2, 4, 5], gamma=0.1)                         # 减小学习速率
 
     """ final options """
     # print(opt)
@@ -266,8 +266,9 @@ def train(opt):
         if i == opt.num_iter:
             print('end the training')
             sys.exit()
-            
-        if i > 0 and i % step_per_epoch == 0:                # 调整学习速率
+
+        if i > 0 and i % int(step_per_epoch) == 0:                # 调整学习速率
+            print('down the learn rate 1/10')
             lrScheduler.step()
 
         i += 1
@@ -283,13 +284,13 @@ def train(opt):
     parser.add_argument('--train_data', default='/home/deepblue/deepbluetwo/chenjun/1_OCR/data/data_lmdb_release/training', help='path to training dataset')
     parser.add_argument('--valid_data', default='/home/deepblue/deepbluetwo/chenjun/1_OCR/data/data_lmdb_release/validation', help='path to validation dataset')
     parser.add_argument('--manualSeed', type=int, default=666, help='for random seed setting')
-    parser.add_argument('--workers', type=int, help='number of data loading workers', default=4)
-    parser.add_argument('--batch_size', type=int, default=16, help='input batch size')
-    parser.add_argument('--num_iter', type=int, default=150000, help='number of iterations to train for')
-    parser.add_argument('--valInterval', type=int, default=50, help='Interval between each validation')
+    parser.add_argument('--workers', type=int, help='number of data loading workers', default=6)
+    parser.add_argument('--batch_size', type=int, default=256, help='input batch size')
+    parser.add_argument('--num_iter', type=int, default=300000, help='number of iterations to train for')
+    parser.add_argument('--valInterval', type=int, default=5000, help='Interval between each validation')
     parser.add_argument('--saveInterval', type=int, default=5000, help='Interval between each save')
     parser.add_argument('--disInterval', type=int, default=5, help='Interval betweet each show')
-    # parser.add_argument('--continue_model', default = '', help="path to model to continue training")
+    parser.add_argument('--continue_model', default = '', help="path to model to continue training")
     # parser.add_argument('--continue_model', default='./saved_models/None-ResNet-SRN-SRN-Seed666/iter_150000.pth', help="path to model to continue training")
     parser.add_argument('--adam', default=True, help='Whether to use adam (default is Adadelta)')
     parser.add_argument('--ranger', default=False, help='use RAdam + Lookahead for optimizer')
@@ -320,7 +321,7 @@ def train(opt):
     parser.add_argument('--batch_max_character', type=int, default=25, help='the max character of one image')
     parser.add_argument('--alphabet_size', type=int, default=None, help='the categry of the string')
 
-    parser.add_argument('--select_data', type=str, default='ICDAR2019-ICDAR2019',
+    parser.add_argument('--select_data', type=str, default='MJ-ST',
                         help='select training data MJ-ST | MJ-ST-ICDAR2019 | baidu')
     parser.add_argument('--batch_ratio', type=str, default='1.0-1.0',
                         help='assign ratio for each selected data in the batch')

utils.py

@@ -175,7 +175,7 @@ def encode(self, text, batch_max_length=25):
         """
         length = [len(s) + 1 for s in text]  # +1 for [s] at end of sentence.
         # additional +1 for [GO] at first step. batch_text is padded with [GO] token after [s] token.
-        batch_text = torch.cuda.LongTensor(len(text), batch_max_length).fill_(self.PAD)
+        batch_text = torch.cuda.LongTensor(len(text), batch_max_length + 1).fill_(self.PAD)
         # mask_text = torch.cuda.LongTensor(len(text), batch_max_length).fill_(0)
         for i, t in enumerate(text):
             t = list(t + self.character[-2])