bceweighted class should only be used if weighting is the intention, else another class is provided

Author: aditya0by0

chebai/loss/bce_weighted.py

@@ -10,7 +10,6 @@
 class BCEWeighted(torch.nn.BCEWithLogitsLoss):
     """
     BCEWithLogitsLoss with weights automatically computed according to the beta parameter.
-    If beta is None or data_extractor is None, the loss is unweighted.
 
     This class computes weights based on the formula from the paper by Cui et al. (2019):
     https://openaccess.thecvf.com/content_CVPR_2019/papers/Cui_Class-Balanced_Loss_Based_on_Effective_Number_of_Samples_CVPR_2019_paper.pdf
@@ -33,19 +32,22 @@ def __init__(
             data_extractor = data_extractor.labeled
         self.data_extractor = data_extractor
 
-        assert (
-            beta is not None
-        ), f"Beta parameter must be provided if this loss ({self.__class__.__name__}) is used."
+        assert self.beta is not None and self.data_extractor is not None, (
+            f"Beta parameter must be provided along with data_extractor, "
+            f"if this loss class ({self.__class__.__name__}) is used."
+        )
 
-        # If beta is provided, require a data_extractor.
-        if self.beta is not None and self.data_extractor is None:
-            raise ValueError("When 'beta' is set, 'data_extractor' must also be set.")
+        assert all(
+            os.path.exists(os.path.join(self.data_extractor.processed_dir, file_name))
+            for file_name in self.data_extractor.processed_file_names
+        ), "Dataset files not found. Make sure the dataset is processed before using this loss."
 
         assert (
             isinstance(self.data_extractor, _ChEBIDataExtractor)
             or self.data_extractor is None
         )
         super().__init__(**kwargs)
+        self.pos_weight: Optional[torch.Tensor] = None
 
     def set_pos_weight(self, input: torch.Tensor) -> None:
         """
@@ -54,17 +56,7 @@ def set_pos_weight(self, input: torch.Tensor) -> None:
         Args:
             input (torch.Tensor): The input tensor for which to set the positive weights.
         """
-        if (
-            self.beta is not None
-            and self.data_extractor is not None
-            and all(
-                os.path.exists(
-                    os.path.join(self.data_extractor.processed_dir, file_name)
-                )
-                for file_name in self.data_extractor.processed_file_names
-            )
-            and self.pos_weight is None
-        ):
+        if self.pos_weight is None:
             print(
                 f"Computing loss-weights based on v{self.data_extractor.chebi_version} dataset (beta={self.beta})"
             )
@@ -105,3 +97,9 @@ def forward(
         """
         self.set_pos_weight(input)
         return super().forward(input, target)
+
+
+class UnWeightedBCEWithLogitsLoss(torch.nn.BCEWithLogitsLoss):
+    def forward(self, input, target, **kwargs):
+        # As the custom passed kwargs are not used in BCEWithLogitsLoss, we can ignore them
+        return super().forward(input, target)