diff --git a/abl/learning/readme.md b/abl/learning/readme.md index 76e7bbc..3d1a003 100644 --- a/abl/learning/readme.md +++ b/abl/learning/readme.md @@ -31,7 +31,7 @@ **num_workers : int** + The number of workers used for loading data. -**criterion : torch.nn.Module** +**loss_fn : torch.nn.Module** + The loss function used for training. **optimizer : torch.nn.Module** @@ -60,13 +60,13 @@ > ```python > # Three necessary component > cls = LeNet5() -> criterion = nn.CrossEntropyLoss() +> loss_fn = nn.CrossEntropyLoss() > optimizer = torch.optim.Adam(cls.parameters()) > > # Initialize base_model > base_model = BasicModel( > cls, -> criterion, +> loss_fn, > optimizer, > torch.device("cuda:0"), > batch_size=32, diff --git a/docs/Intro/Learning.rst b/docs/Intro/Learning.rst index 8c4504f..66d644b 100644 --- a/docs/Intro/Learning.rst +++ b/docs/Intro/Learning.rst @@ -39,11 +39,11 @@ For a PyTorch-based neural network, we first need to encapsulate it within a ``B import torchvision cls = torchvision.models.resnet18(pretrained=True) - # criterion and optimizer are used for training - criterion = torch.nn.CrossEntropyLoss() + # loss_fn and optimizer are used for training + loss_fn = torch.nn.CrossEntropyLoss() optimizer = torch.optim.Adam(cls.parameters()) - base_model = BasicNN(cls, criterion, optimizer) + base_model = BasicNN(cls, loss_fn, optimizer) model = ABLModel(base_model) Besides ``fit`` and ``predict``, ``BasicNN`` also implements the following methods: diff --git a/docs/Intro/Quick-Start.rst b/docs/Intro/Quick-Start.rst index 9f7f9f1..f5d4cf6 100644 --- a/docs/Intro/Quick-Start.rst +++ b/docs/Intro/Quick-Start.rst @@ -84,17 +84,17 @@ To build the machine learning part, we need to wrap our machine learning model i # The number of pseudo labels is 10 cls = LeNet5(num_classes=10) -Aside from the network, we need to define a criterion, an optimizer, and a device so as to create a ``BasicNN`` object. This class implements ``fit``, ``predict``, ``predict_proba`` and several other methods to enable the PyTorch-based neural network to work as a scikit-learn model. +Aside from the network, we need to define a loss_fn, an optimizer, and a device so as to create a ``BasicNN`` object. This class implements ``fit``, ``predict``, ``predict_proba`` and several other methods to enable the PyTorch-based neural network to work as a scikit-learn model. .. code:: python import torch from abl.learning import BasicNN - criterion = torch.nn.CrossEntropyLoss() + loss_fn = torch.nn.CrossEntropyLoss() optimizer = torch.optim.Adam(cls.parameters(), lr=0.001, betas=(0.9, 0.99)) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - base_model = BasicNN(cls, criterion, optimizer, device) + base_model = BasicNN(cls, loss_fn, optimizer, device) .. code:: python diff --git a/examples/hed/hed_bridge.py b/examples/hed/hed_bridge.py index 0b08ec8..8543d1f 100644 --- a/examples/hed/hed_bridge.py +++ b/examples/hed/hed_bridge.py @@ -32,14 +32,14 @@ class HEDBridge(SimpleBridge): num_classes=len(self.reasoner.kb.pseudo_label_list) ) device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") - criterion = torch.nn.MSELoss() + loss_fn = torch.nn.MSELoss() optimizer = torch.optim.RMSprop( cls_autoencoder.parameters(), lr=0.001, alpha=0.9, weight_decay=1e-6 ) pretrain_model = BasicNN( cls_autoencoder, - criterion, + loss_fn, optimizer, device, save_interval=1, diff --git a/examples/hed/hed_example.ipynb b/examples/hed/hed_example.ipynb index 0e9c9d4..7858727 100644 --- a/examples/hed/hed_example.ipynb +++ b/examples/hed/hed_example.ipynb @@ -157,7 +157,7 @@ "source": [ "# Build necessary components for BasicNN\n", "cls = SymbolNet(num_classes=4)\n", - "criterion = nn.CrossEntropyLoss()\n", + "loss_fn = nn.CrossEntropyLoss()\n", "optimizer = torch.optim.RMSprop(cls.parameters(), lr=0.001, weight_decay=1e-6)\n", "# optimizer = torch.optim.Adam(cls.parameters(), lr=0.001, betas=(0.9, 0.99))\n", "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")" @@ -173,7 +173,7 @@ "# The function of BasicNN is to wrap NN models into the form of an sklearn estimator\n", "base_model = BasicNN(\n", " cls,\n", - " criterion,\n", + " loss_fn,\n", " optimizer,\n", " device,\n", " batch_size=32,\n", diff --git a/examples/hwf/hwf_example.ipynb b/examples/hwf/hwf_example.ipynb index bae8d0d..0a31e18 100644 --- a/examples/hwf/hwf_example.ipynb +++ b/examples/hwf/hwf_example.ipynb @@ -95,7 +95,7 @@ "# Initialize necessary component for machine learning part\n", "cls = SymbolNet(num_classes=len(kb.pseudo_label_list), image_size=(45, 45, 1))\n", "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n", - "criterion = nn.CrossEntropyLoss()\n", + "loss_fn = nn.CrossEntropyLoss()\n", "optimizer = torch.optim.Adam(cls.parameters(), lr=0.001, betas=(0.9, 0.99))" ] }, @@ -109,7 +109,7 @@ "# The function of BasicNN is to wrap NN models into the form of an sklearn estimator\n", "base_model = BasicNN(\n", " model=cls,\n", - " criterion=criterion,\n", + " loss_fn=loss_fn,\n", " optimizer=optimizer,\n", " device=device,\n", " save_interval=1,\n", @@ -199,7 +199,7 @@ "metadata": {}, "outputs": [], "source": [ - "bridge.train(train_data, loops=3, segment_size=1000, save_interval=1, save_dir=weights_dir)\n", + "bridge.train(train_data, train_data, loops=3, segment_size=1000, save_interval=1, save_dir=weights_dir)\n", "bridge.test(test_data)" ] } diff --git a/examples/mnist_add/mnist_add_example.ipynb b/examples/mnist_add/mnist_add_example.ipynb index 980094f..53d5bce 100644 --- a/examples/mnist_add/mnist_add_example.ipynb +++ b/examples/mnist_add/mnist_add_example.ipynb @@ -50,7 +50,7 @@ "source": [ "# Build necessary components for BasicNN\n", "cls = LeNet5(num_classes=10)\n", - "criterion = nn.CrossEntropyLoss()\n", + "loss_fn = nn.CrossEntropyLoss()\n", "optimizer = torch.optim.Adam(cls.parameters(), lr=0.001, betas=(0.9, 0.99))\n", "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")" ] @@ -65,7 +65,7 @@ "# The function of BasicNN is to wrap NN models into the form of an sklearn estimator\n", "base_model = BasicNN(\n", " cls,\n", - " criterion,\n", + " loss_fn,\n", " optimizer,\n", " device,\n", " batch_size=32,\n", diff --git a/examples/models/nn.py b/examples/models/nn.py index 8720b36..5fb7f3e 100644 --- a/examples/models/nn.py +++ b/examples/models/nn.py @@ -66,7 +66,7 @@ class SymbolNet(nn.Module): num_features = 64 * (image_size[0] // 4 - 1) * (image_size[1] // 4 - 1) self.fc1 = nn.Sequential(nn.Linear(num_features, 120), nn.ReLU()) self.fc2 = nn.Sequential(nn.Linear(120, 84), nn.ReLU()) - self.fc3 = nn.Sequential(nn.Linear(84, num_classes), nn.Softmax(dim=1)) + self.fc3 = nn.Sequential(nn.Linear(84, num_classes)) def forward(self, x): x = self.conv1(x) diff --git a/tests/conftest.py b/tests/conftest.py index 2590308..73e7029 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -44,18 +44,18 @@ class LeNet5(nn.Module): @pytest.fixture def basic_nn_instance(): model = LeNet5() - criterion = nn.CrossEntropyLoss() + loss_fn = nn.CrossEntropyLoss() optimizer = optim.Adam(model.parameters()) - return BasicNN(model, criterion, optimizer) + return BasicNN(model, loss_fn, optimizer) # Fixture for base_model instance @pytest.fixture def base_model_instance(): model = LeNet5() - criterion = nn.CrossEntropyLoss() + loss_fn = nn.CrossEntropyLoss() optimizer = optim.Adam(model.parameters()) - return BasicNN(model, criterion, optimizer) + return BasicNN(model, loss_fn, optimizer) # Fixture for ListData instance diff --git a/tests/test_basic_nn.py b/tests/test_basic_nn.py index 0a6cafb..b0bdada 100644 --- a/tests/test_basic_nn.py +++ b/tests/test_basic_nn.py @@ -27,7 +27,7 @@ class TestBasicNN(object): def test_initialization(self, basic_nn_instance): """Test initialization of the BasicNN class""" assert basic_nn_instance.model is not None - assert isinstance(basic_nn_instance.criterion, nn.Module) + assert isinstance(basic_nn_instance.loss_fn, nn.Module) assert isinstance(basic_nn_instance.optimizer, optim.Optimizer) def test_training_methods(self, basic_nn_instance, sample_data, sample_data_loader_with_label):