beta0.1

app.py

@@ -24,7 +24,7 @@ class App:
 
     def train(self, project_name: str):
         logger.info("\nStart Train ----> {}\n".format(project_name))
-        trainer = train.Train("test1")
+        trainer = train.Train(project_name)
         trainer.start()
 
 

configs/base.py

@@ -28,7 +28,7 @@ class Config(object):
                 "BATCH_SIZE": 32,
                 "TEST_BATCH_SIZE": 32,
                 'CNN': {
-                    "NAME": "ddddOcr",
+                    "NAME": "ddddocr",
                 },
                 'DROPOUT': 0.3,
                 'OPTIMIZER': 'SGD',

nets/__init__.py

@@ -2,6 +2,7 @@ import json
 
 from .backbone import *
 import torch
+torch.set_num_threads(1)
 
 
 class Net(torch.nn.Module):
@@ -46,8 +47,7 @@ class Net(torch.nn.Module):
         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.lstm = torch.nn.LSTM(input_size=self.out_size, hidden_size=self.out_size, bidirectional=True, num_layers=1, dropout=self.dropout)
             self.paramters.append({'params': self.lstm.parameters()})
 
             self.loss = torch.nn.CTCLoss(blank=0, reduction='mean')
@@ -97,13 +97,13 @@ class Net(torch.nn.Module):
             outputs = self.fc(outputs)
         return outputs
 
-    def training(self, inputs, labels, labels_length):
+    def trainer(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)
+    def tester(self, inputs, labels, labels_length):
+        predict = self.get_features(inputs)
         pred_decode_labels = []
         labels_list = []
         correct_list = []
@@ -136,7 +136,7 @@ class Net(torch.nn.Module):
             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():
+            if labels_list[ids] == pred_decode_labels[ids]:
                 correct_list.append(ids)
             else:
                 error_list.append(ids)
@@ -147,7 +147,7 @@ class Net(torch.nn.Module):
         if self.word:
             loss = self.loss(predict, labels.long().cuda())
         else:
-            log_predict = predict.log_softmax(2).detach().requires_grad_()
+            log_predict = predict.log_softmax(2)
             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()

utils/cache_data.py

@@ -63,8 +63,7 @@ class CacheData:
     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')
@@ -88,6 +87,8 @@ class CacheData:
             else:
                 logger.warning("\nFile({}) has a suffix that is not allowed! We will remove it!".format(file))
         labels = list(set(labels))
+        if not self.conf['Model']['Word']:
+            labels.insert(0, " ")
         logger.info("\nCoolect labels is {}".format(json.dumps(labels, ensure_ascii=False)))
         self.conf['System']['Path'] = base_path
         self.conf['Model']['CharSet'] = labels

utils/train.py

@@ -64,7 +64,7 @@ class 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)
+                loss, lr = self.net.trainer(inputs, labels, labels_length)
 
                 self.avg_loss += loss
 
@@ -80,9 +80,11 @@ class Train:
                     model_path = os.path.join(self.checkpoints_path, "checkpoint_{}_{}_{}.tar".format(
                         self.project_name, self.epoch, self.step,
                     ))
+                    self.net.scheduler.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)
@@ -94,7 +96,7 @@ class Train:
                         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,
+                    pred_labels, labels_list, correct_list, error_list = self.net.tester(test_inputs, test_labels,
                                                                                           test_labels_length)
                     self.net = self.net.train()
                     accuracy = len(correct_list) / test_inputs.shape[0]
@@ -102,6 +104,7 @@ class Train:
                         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
                     ))
+                    self.avg_loss = 0
                     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()
@@ -121,7 +124,7 @@ class Train:
                         exit()
 
             self.epoch += 1
-            self.net.scheduler.step(self.epoch)
+
 
 
 if __name__ == '__main__':