Fix the error of no num_class (#2200)

PaddlePaddle · Jun 10, 2022 · 00080d9 · 00080d9
1 parent d9da1d2
commit 00080d9
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/paddleseg/cvlibs/config.py b/paddleseg/cvlibs/config.py
@@ -315,17 +315,22 @@ def model(self) -> paddle.nn.Layer:
         model_cfg = self.dic.get('model').copy()
         if not model_cfg:
             raise RuntimeError('No model specified in the configuration file.')
+
         if not 'num_classes' in model_cfg:
             num_classes = None
             try:
                 if self.train_dataset_config:
                     if hasattr(self.train_dataset_class, 'NUM_CLASSES'):
                         num_classes = self.train_dataset_class.NUM_CLASSES
+                    elif 'num_classes' in self.train_dataset_config:
+                        num_classes = self.train_dataset_config['num_classes']
                     elif hasattr(self.train_dataset, 'num_classes'):
                         num_classes = self.train_dataset.num_classes
                 elif self.val_dataset_config:
                     if hasattr(self.val_dataset_class, 'NUM_CLASSES'):
                         num_classes = self.val_dataset_class.NUM_CLASSES
+                    elif 'num_classes' in self.val_dataset_config:
+                        num_classes = self.val_dataset_config['num_classes']
                     elif hasattr(self.val_dataset, 'num_classes'):
                         num_classes = self.val_dataset.num_classes
             except FileNotFoundError: