From 2d54e7023d45860f93041ff6f7f42028c4882c6f Mon Sep 17 00:00:00 2001 From: troyyyyy Date: Thu, 16 Mar 2023 09:49:44 +0800 Subject: [PATCH] Update HED example --- abl/framework_hed.py | 6 ++---- examples/example.py | 42 +++++++++++++++++++++--------------------- 2 files changed, 23 insertions(+), 25 deletions(-) diff --git a/abl/framework_hed.py b/abl/framework_hed.py index ab88318..76c57fa 100644 --- a/abl/framework_hed.py +++ b/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))) diff --git a/examples/example.py b/examples/example.py index 39e8c3f..0f75895 100644 --- a/examples/example.py +++ b/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