refine infer

lwplwp123 · Jan 28, 2019 · 4e140fa · 4e140fa
1 parent fa4f838
commit 4e140fa
Show file tree

Hide file tree

Showing 4 changed files with 33 additions and 9 deletions.
diff --git a/detector/ctpn/infer.py b/detector/ctpn/infer.py
@@ -342,8 +342,12 @@ def random_test(net):
 
 if __name__ == '__main__':
     running_mode = sys.argv[2]  # cpu or gpu
+    print("Mode: %s" % running_mode)
     net = Net.CTPN()
-    net.load_state_dict(torch.load(MODEL, map_location=running_mode))
+    if running_mode == 'cpu':
+        net.load_state_dict(torch.load(MODEL, map_location=running_mode))
+    else:
+        net.load_state_dict(torch.load(MODEL))
     print(net)
     net.eval()
 

diff --git a/recognizer/crnn/Config.py b/recognizer/crnn/Config.py
@@ -8,7 +8,7 @@
 gpu_id = '4'
 model_dir = './model'
 data_worker = 4
-batch_size = 128
+batch_size = 64
 img_height = 32
 img_width = 100
 alphabet = alphabets.alphabet