File size: 1,874 Bytes
bc45d7d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 | from dataclasses import dataclass, field
import transformers
import timm
@dataclass
class LstmConfig:
input_size: int = 134
hidden_size: int = 256
num_layers: int = 5
batch_first: bool = True
bidirectional: bool = True
dropout: float = 0.2
@dataclass
class XgbConfig:
booster: str = "gbtree"
silent: int = 0
max_depth: int = 2
subsample: float = 0.9923301318585108
colsample_bytree: float = 0.7747027267489391
reg_lambda: int = 3
objective: str = "multi:softprob"
eval_metric: str = "mlogloss"
tree_method: str = "gpu_hist" ## change it to `hist` if gpu not available
@dataclass
class TransformerConfig:
size: str
input_size: int = 134
max_position_embeddings: int = field(default=256, repr=False)
layer_norm_eps: float = field(default=1e-12, repr=False)
hidden_dropout_prob: float = field(default=0.1, repr=False)
hidden_size: int = field(default=512, repr=False)
num_attention_heads: int = field(default=8, repr=False)
num_hidden_layers: int = field(default=4, repr=False)
model_config: transformers.BertConfig = field(init=False)
def __post_init__(self):
assert self.size in ["small", "large"]
if self.size == "small":
self.hidden_size = 256
self.num_attention_heads = 4
self.num_hidden_layers = 2
self.model_config = transformers.BertConfig(
hidden_size=self.hidden_size,
num_attention_heads=self.num_attention_heads,
num_hidden_layers=self.num_hidden_layers,
max_position_embeddings=self.max_position_embeddings,
)
@dataclass
class CnnConfig:
model: str = "mobilenetv2_100"
output_dim: int = 1280
def __post_init__(self):
available_models = timm.list_models(pretrained=True)
assert self.model in available_models
|