File size: 2,221 Bytes
1197f7d
b5fa3f1
1197f7d
856cce6
 
1197f7d
cbbfcfe
 
 
 
 
 
b5fa3f1
 
 
 
 
 
 
 
 
 
 
 
1197f7d
 
cbbfcfe
b5fa3f1
1197f7d
 
 
b5fa3f1
 
 
 
 
 
 
 
 
 
 
 
1197f7d
b5fa3f1
1197f7d
 
 
b5fa3f1
1197f7d
b5fa3f1
1197f7d
b5fa3f1
1197f7d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
b5fa3f1
 
 
 
 
 
 
 
 
 
 
1197f7d
 
 
 
 
b5fa3f1
 
1197f7d
 
 
 
 
 
 
 
 
 
b5fa3f1
 
6e85a96
b5fa3f1
1197f7d
b5fa3f1
1197f7d
 
 
 
16c6705
b5fa3f1
 
 
16c6705
 
1197f7d
b5fa3f1
 
 
1197f7d
 
 
b5fa3f1
 
 
 
1197f7d
b5fa3f1
 
1197f7d
b5fa3f1
 
1197f7d
b5fa3f1
 
1197f7d
b5fa3f1
 
 
1197f7d
306fc38
 
1197f7d
856cce6
 
 
 
 
 
5727efb
856cce6
 
 
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
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
from dataclasses import dataclass
from typing import Any, Dict, List, Optional, Union

from torch import nn


@dataclass
class AnchorConfig:
    reg_max: int
    strides: List[int]


@dataclass
class LayerConfg:
    args: Dict
    source: Union[int, str, List[int]]
    tags: str


@dataclass
class BlockConfig:
    block: List[Dict[str, LayerConfg]]


@dataclass
class Model:
    anchor: AnchorConfig
    model: Dict[str, BlockConfig]


@dataclass
class DownloadDetail:
    url: str
    file_size: int


@dataclass
class DownloadOptions:
    details: Dict[str, DownloadDetail]


@dataclass
class DatasetConfig:
    path: str
    auto_download: Optional[DownloadOptions]


@dataclass
class DataConfig:
    shuffle: bool
    batch_size: int
    pin_memory: bool
    data_augment: Dict[str, int]


@dataclass
class OptimizerArgs:
    lr: float
    weight_decay: float


@dataclass
class OptimizerConfig:
    type: str
    args: OptimizerArgs


@dataclass
class MatcherConfig:
    iou: str
    topk: int
    factor: Dict[str, int]


@dataclass
class LossConfig:
    objective: Dict[str, int]
    aux: Union[bool, float]
    matcher: MatcherConfig


@dataclass
class SchedulerConfig:
    type: str
    warmup: Dict[str, Union[int, float]]
    args: Dict[str, Any]


@dataclass
class EMAConfig:
    enabled: bool
    decay: float


@dataclass
class TrainConfig:
    task: str
    dataset: DatasetConfig
    epoch: int
    data: DataConfig
    optimizer: OptimizerConfig
    loss: LossConfig
    scheduler: SchedulerConfig
    ema: EMAConfig


@dataclass
class NMSConfig:
    min_confidence: int
    min_iou: int


@dataclass
class InferenceConfig:
    task: str
    nms: NMSConfig


@dataclass
class Config:
    task: Union[TrainConfig, InferenceConfig]
    model: Model
    name: str

    device: Union[str, int, List[int]]
    cpu_num: int

    class_num: int
    image_size: List[int]

    out_path: str
    exist_ok: bool

    lucky_number: 10
    use_wandb: bool
    use_TensorBoard: bool

    weight: Optional[str]


@dataclass
class YOLOLayer(nn.Module):
    source: Union[int, str, List[int]]
    output: bool
    tags: str
    layer_type: str
    usable: bool

    def __post_init__(self):
        super().__init__()