add weight fusion and class train

Author: zlaom

run2.sh

@@ -94,13 +94,26 @@ fi
 # #################### execute according to TYPE
 ########## train
 if [ "$TYPE" = "train" ]; then
-    cd src/text_predict/ocr_bert_base
-    time python ocr_classifier.py
+#     cd src/text_predict/ocr_bert_base
+#     time python ocr_classifier.py
+#     cd -
+    cd src/text_predict/ocr_bert_class_train
+#     time python train.py
+    cd -
+    
+    cd src/weight_fusion
+    time python train.py
     cd -
   exit 0
-fi
-########## test
-# elif [ "$TYPE" = "test" ]; then
 
-#   exit 0
+########## test
+elif [ "$TYPE" = "test" ]; then
+    cd src/text_predict/ocr_bert_class_train
+    time python predict.py
+    cd -
+    cd src/weight_fusion
+    time python predict.py
+    cd -
+  exit 0
 # ######### text predict
+fi

src/test_2nd.txt

@@ -134,7 +134,6 @@ def train(run_i, model, loss_fn, config,  train_loader, val_loader, root_path):
                 num+=1
                 loss.backward()
                 optimizer.step()
-                break
         gap = evaluate(config, val_loader, tokenizer, model)
         with open(root_path+'res.csv', 'a+') as out:
             writer = csv.writer(out)
@@ -154,7 +153,6 @@ def train(run_i, model, loss_fn, config,  train_loader, val_loader, root_path):
         if(count>3):
             print("overfit stop train!")
             break
-        break
 
 # pre process
 def predict(config, model,checkpoint,predict_loader,reslut_path):