darknet_classification.py 2.51 KB
Newer Older
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
1
# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
2
3
4
5
6
7
8
9
10
11
12
13
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
Yeqing Li's avatar
Yeqing Li committed
14

15
16
"""Image classification with darknet configs."""

17
import dataclasses
Abdullah Rashwan's avatar
Abdullah Rashwan committed
18
from typing import List, Optional
19

20
21
22
from official.core import config_definitions as cfg
from official.core import exp_factory
from official.modeling import hyperparams
Abdullah Rashwan's avatar
Abdullah Rashwan committed
23
from official.projects.yolo.configs import backbones
Abdullah Rashwan's avatar
Abdullah Rashwan committed
24
25
from official.vision.configs import common
from official.vision.configs import image_classification as imc
26
27
28
29


@dataclasses.dataclass
class ImageClassificationModel(hyperparams.Config):
30
  """Image classification model config."""
31
  num_classes: int = 0
32
  input_size: List[int] = dataclasses.field(default_factory=lambda: [224, 224])
33
  backbone: backbones.Backbone = backbones.Backbone(
Jaeyoun Kim's avatar
Jaeyoun Kim committed
34
      type='darknet', darknet=backbones.Darknet())
35
36
  dropout_rate: float = 0.0
  norm_activation: common.NormActivation = common.NormActivation()
Abdullah Rashwan's avatar
Abdullah Rashwan committed
37
  # Adds a Batch Normalization layer pre-GlobalAveragePooling in classification.
38
  add_head_batch_norm: bool = False
Vishnu Banna's avatar
Vishnu Banna committed
39
  kernel_initializer: str = 'VarianceScaling'
40
41
42
43
44
45
46
47


@dataclasses.dataclass
class Losses(hyperparams.Config):
  one_hot: bool = True
  label_smoothing: float = 0.0
  l2_weight_decay: float = 0.0

48

49
50
51
52
53
54
@dataclasses.dataclass
class ImageClassificationTask(cfg.TaskConfig):
  """The model config."""
  model: ImageClassificationModel = ImageClassificationModel()
  train_data: imc.DataConfig = imc.DataConfig(is_training=True)
  validation_data: imc.DataConfig = imc.DataConfig(is_training=False)
55
  evaluation: imc.Evaluation = imc.Evaluation()
56
57
  losses: Losses = Losses()
  gradient_clip_norm: float = 0.0
58
  logging_dir: Optional[str] = None
59
60
61


@exp_factory.register_config_factory('darknet_classification')
62
def darknet_classification() -> cfg.ExperimentConfig:
63
64
65
66
67
68
69
  """Image classification general."""
  return cfg.ExperimentConfig(
      task=ImageClassificationTask(),
      trainer=cfg.TrainerConfig(),
      restrictions=[
          'task.train_data.is_training != None',
          'task.validation_data.is_training != None'
anivegesana's avatar
anivegesana committed
70
      ])