Update HED example

abl/framework_hed.py

@@ -172,9 +172,6 @@ def abduce_and_train(model, abducer, mapping, train_X_true, select_num):
     if len(consistent_idx) == 0:
         return 0, 0, None
     
-    if len(mappings) > 1:
-        INFO('Final mapping is: ', mapping)
-    
     INFO('Train pool size is:', len(flatten(consistent_pred_res)))
     INFO("Start to use abduced pseudo label to train model...")
     model.train([X[idx] for idx in consistent_idx], remapping_res(consistent_pred_res, mapping))
@@ -214,7 +211,7 @@ def get_rules_from_data(model, abducer, mapping, train_X_true, samples_per_rule,
             consistent_idx = []
             consistent_pred_res = []
             for idx in range(len(pred_res)):
-                if abducer.kb.logic_forward(pred_res[idx]):
+                if abducer.kb.logic_forward([pred_res[idx]]):
                     consistent_idx.append(idx)
                     consistent_pred_res.append(pred_res[idx])
 
@@ -296,6 +293,7 @@ def train_with_rule(model, abducer, train_data, val_data, select_num=10, min_len
                     break
                 else:
                     if equation_len == min_len:
+                        INFO('Final mapping is: ', mapping)
                         model.cls_list[0].model.load_state_dict(torch.load("./weights/pretrain_weights.pth"))
                     else:
                         model.cls_list[0].model.load_state_dict(torch.load("./weights/weights_%d.pth" % (equation_len - 1)))

examples/example.py

@@ -22,8 +22,8 @@ from abl.models.basic_model import BasicModel, BasicDataset
 from abl.models.wabl_models import DecisionTree, WABLBasicModel
 
 from multiprocessing import Pool
-from abl.abducer.abducer_base import AbducerBase
-from abl.abducer.kb import add_KB, HWF_KB, prolog_KB
+from abl.abducer.abducer_base import AbducerBase, HED_Abducer
+from abl.abducer.kb import add_KB, HWF_KB, prolog_KB, HED_prolog_KB
 from datasets.mnist_add.get_mnist_add import get_mnist_add
 from datasets.hwf.get_hwf import get_hwf
 from datasets.hed.get_hed import get_hed, split_equation
@@ -32,37 +32,37 @@ from abl import framework_hed
 
 def run_test():
 
-    kb = add_KB()
-    # kb = HWF_KB(True)
-    abducer = AbducerBase(kb, 'confidence')
+    # kb = add_KB()
+    # kb = HWF_KB()
+    # abducer = AbducerBase(kb, 'confidence')
 
-    # kb = prolog_KB(pseudo_label_list=[1, 0, '+', '='], pl_file='../examples/datasets/hed/learn_add.pl')
-    # abducer = AbducerBase(kb, zoopt=True, multiple_predictions=True)
+    kb = HED_prolog_KB(pseudo_label_list=[1, 0, '+', '='], pl_file='../examples/datasets/hed/learn_add.pl')
+    abducer = HED_Abducer(kb)
 
     recorder = logger()
 
-    # total_train_data = get_hed(train=True)
-    # train_data, val_data = split_equation(total_train_data, 3, 1)
-    # test_data = get_hed(train=False)
+    total_train_data = get_hed(train=True)
+    train_data, val_data = split_equation(total_train_data, 3, 1)
+    test_data = get_hed(train=False)
     
-    train_data = get_mnist_add(train = True, get_pseudo_label = True)
-    test_data = get_mnist_add(train = False, get_pseudo_label = True)
+    # train_data = get_mnist_add(train = True, get_pseudo_label = True)
+    # test_data = get_mnist_add(train = False, get_pseudo_label = True)
 
     # train_data = get_hwf(train = True, get_pseudo_label = True)
     # test_data = get_hwf(train = False, get_pseudo_label = True)
     
-    cls = LeNet5(num_classes=len(kb.pseudo_label_list), image_size=(train_data[0][0][0].shape[1:]))
-    # cls = SymbolNet(num_classes=len(kb.pseudo_label_list))
+    # cls = LeNet5(num_classes=len(kb.pseudo_label_list), image_size=(train_data[0][0][0].shape[1:]))
+    cls = SymbolNet(num_classes=len(kb.pseudo_label_list))
     device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
     
-    # framework_hed.hed_pretrain(kb, cls, recorder)
+    framework_hed.hed_pretrain(kb, cls, recorder)
     
     criterion = nn.CrossEntropyLoss()
-    # optimizer = torch.optim.RMSprop(cls.parameters(), lr=0.001, weight_decay=1e-6)
-    optimizer = torch.optim.Adam(cls.parameters(), lr=0.001, betas=(0.9, 0.99))
+    optimizer = torch.optim.RMSprop(cls.parameters(), lr=0.001, weight_decay=1e-6)
+    # optimizer = torch.optim.Adam(cls.parameters(), lr=0.001, betas=(0.9, 0.99))
     
 
-    base_model = BasicModel(cls, criterion, optimizer, device, save_interval=1, save_dir=recorder.save_dir, batch_size=32, num_epochs=1, recorder=recorder)
+    base_model = BasicModel(cls, criterion, optimizer, device, save_interval=1, save_dir=recorder.save_dir, batch_size=32, num_epochs=10, recorder=recorder)
     model = WABLBasicModel(base_model, kb.pseudo_label_list)
     
     # train_X, train_Z, train_Y = get_mnist_add(train = True, get_pseudo_label = True)
@@ -71,10 +71,10 @@ def run_test():
     # train_data = get_hwf(train = True, get_pseudo_label = True)
     # test_data = get_hwf(train = False, get_pseudo_label = True)
 
-    # model, mapping = framework_hed.train_with_rule(model, abducer, train_data, val_data, select_num=10, min_len=5, max_len=8)
-    # framework_hed.hed_test(model, abducer, mapping, train_data, test_data, min_len=5, max_len=8)
+    model, mapping = framework_hed.train_with_rule(model, abducer, train_data, val_data, select_num=10, min_len=5, max_len=8)
+    framework_hed.hed_test(model, abducer, mapping, train_data, test_data, min_len=5, max_len=8)
     
-    framework_hed.train(model, abducer, train_data, test_data, sample_num=10000, verbose=1)
+    # framework_hed.train(model, abducer, train_data, test_data, sample_num=10000, verbose=1)
 
     recorder.dump()
     return True