File size: 394 Bytes
2e10f4f
 
 
 
 
f940030
2e10f4f
 
 
f940030
1
2
3
4
5
6
7
8
9
10
from transformers import PretrainedConfig

class SqueezeNetConfig(PretrainedConfig):
    model_type = "squeezenet_custom"

    def __init__(self, num_classes: int = 200, fire_dropout: float = 0.0, final_dropout: float = 0.50, **kwargs):
        self.num_classes = num_classes
        self.fire_dropout = fire_dropout
        self.final_dropout = final_dropout
        super().__init__(**kwargs)