Selaa lähdekoodia

Merge pull request #1295 from will-jl944/develop_jf

disable picodet pruning
will-jl944 4 vuotta sitten
vanhempi
commit
9cd56cb1f9
1 muutettua tiedostoa jossa 2 lisäystä ja 2 poistoa
  1. 2 2
      paddlex/cv/models/base.py

+ 2 - 2
paddlex/cv/models/base.py

@@ -444,11 +444,11 @@ class BaseModel:
             criterion({'l1_norm', 'fpgm'}, optional): Pruning criterion. Defaults to 'l1_norm'.
             save_dir(str, optional): The directory to save sensitivity file of the model. Defaults to 'output'.
         """
-        if self.__class__.__name__ in ['FasterRCNN', 'MaskRCNN']:
+        if self.__class__.__name__ in {'FasterRCNN', 'MaskRCNN', 'PicoDet'}:
             raise Exception("{} does not support pruning currently!".format(
                 self.__class__.__name__))
 
-        assert criterion in ['l1_norm', 'fpgm'], \
+        assert criterion in {'l1_norm', 'fpgm'}, \
             "Pruning criterion {} is not supported. Please choose from ['l1_norm', 'fpgm']"
         arrange_transforms(
             model_type=self.model_type,