diff --git a/example.py b/example.py index dad33f7..7ceaeff 100644 --- a/example.py +++ b/example.py @@ -15,7 +15,7 @@ import framework import torch.nn as nn import torch -from models.lenet5 import LeNet5 +from models.lenet5 import LeNet5, SymbolNet from models.basic_model import BasicModel from models.wabl_models import WABLBasicModel @@ -28,16 +28,20 @@ from datasets.hwf.get_hwf import get_hwf def run_test(): - kb = add_KB() - # kb = hwf_KB() + # kb = add_KB(True) + kb = hwf_KB(True) abducer = AbducerBase(kb) recorder = logger() - train_X, train_Z, train_Y = get_mnist_add(train = True, get_pseudo_label = True) - test_X, test_Z, test_Y = get_mnist_add(train = False, get_pseudo_label = True) + # train_X, train_Z, train_Y = get_mnist_add(train = True, get_pseudo_label = True) + # test_X, test_Z, test_Y = 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_X[0][0].shape[1:])) + # 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)) criterion = nn.CrossEntropyLoss() optimizer = torch.optim.Adam(cls.parameters(), lr=0.001, betas=(0.9, 0.99)) @@ -46,12 +50,11 @@ def run_test(): base_model = BasicModel(cls, criterion, optimizer, device, save_interval=1, save_dir=recorder.save_dir, num_epochs=1, recorder=recorder) model = WABLBasicModel(base_model, kb.pseudo_label_list) - res = framework.train(model, abducer, train_X, train_Z, train_Y, sample_num = 10000, verbose = 1) - recorder.print("abl_acc is ", res) + res = framework.train(model, abducer, train_data, test_data, sample_num = 10000, verbose = 1) + recorder.print(res) recorder.dump() return True if __name__ == "__main__": run_test() -