File size: 1,672 Bytes
1197f7d 856cce6 1197f7d cbbfcfe 1197f7d cbbfcfe 1197f7d cbbfcfe 1197f7d cbbfcfe 1197f7d cbbfcfe 1197f7d 856cce6 1197f7d |
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 |
from dataclasses import dataclass
from typing import Dict, List, Union
from torch import nn
@dataclass
class AnchorConfig:
reg_max: int
strides: List[int]
@dataclass
class Model:
anchor: AnchorConfig
model: Dict[str, List[Dict[str, Union[Dict, List, int]]]]
@dataclass
class Download:
auto: bool
path: str
@dataclass
class DataLoaderConfig:
batch_size: int
shuffle: bool
num_workers: int
pin_memory: bool
image_size: List[int]
class_num: int
@dataclass
class OptimizerArgs:
lr: float
weight_decay: float
@dataclass
class OptimizerConfig:
type: str
args: OptimizerArgs
@dataclass
class SchedulerArgs:
step_size: int
gamma: float
@dataclass
class SchedulerConfig:
type: str
args: SchedulerArgs
@dataclass
class EMAConfig:
enabled: bool
decay: float
@dataclass
class MatcherConfig:
iou: str
topk: int
factor: Dict[str, int]
@dataclass
class TrainConfig:
optimizer: OptimizerConfig
scheduler: SchedulerConfig
ema: EMAConfig
matcher: MatcherConfig
@dataclass
class HyperConfig:
data: DataLoaderConfig
train: TrainConfig
@dataclass
class Dataset:
file_name: str
num_files: int
@dataclass
class Datasets:
base_url: str
images: Dict[str, Dataset]
@dataclass
class Download:
auto: bool
save_path: str
datasets: Datasets
@dataclass
class YOLOLayer(nn.Module):
source: Union[int, str, List[int]]
output: bool
tags: str
layer_type: str
def __post_init__(self):
super().__init__()
@dataclass
class Config:
model: Model
download: Download
hyper: HyperConfig
|