basnet.py 5.33 KB
Newer Older
1
# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
Gunho Park's avatar
Gunho Park committed
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.
14

Gunho Park's avatar
Gunho Park committed
15
16
"""BASNet configuration definition."""
import dataclasses
17
18
import os
from typing import List, Optional, Union
19
from official.core import config_definitions as cfg
Gunho Park's avatar
Gunho Park committed
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
from official.core import exp_factory
from official.modeling import hyperparams
from official.modeling import optimization
from official.vision.beta.configs import common


@dataclasses.dataclass
class DataConfig(cfg.DataConfig):
  """Input config for training."""
  output_size: List[int] = dataclasses.field(default_factory=list)
  # If crop_size is specified, image will be resized first to
  # output_size, then crop of size crop_size will be cropped.
  crop_size: List[int] = dataclasses.field(default_factory=list)
  input_path: str = ''
  global_batch_size: int = 0
  is_training: bool = True
  dtype: str = 'float32'
  shuffle_buffer_size: int = 1000
  cycle_length: int = 10
  resize_eval_groundtruth: bool = True
  groundtruth_padded_size: List[int] = dataclasses.field(default_factory=list)
  aug_rand_hflip: bool = True
42
  file_type: str = 'tfrecord'
Gunho Park's avatar
Gunho Park committed
43
44
45
46
47
48


@dataclasses.dataclass
class BASNetModel(hyperparams.Config):
  """BASNet model config."""
  input_size: List[int] = dataclasses.field(default_factory=list)
Gunho Park's avatar
Gunho Park committed
49
  use_bias: bool = False
Gunho Park's avatar
Gunho Park committed
50
51
52
53
54
55
  norm_activation: common.NormActivation = common.NormActivation()


@dataclasses.dataclass
class Losses(hyperparams.Config):
  label_smoothing: float = 0.1
56
  ignore_label: int = 0  # will be treated as background
Gunho Park's avatar
Gunho Park committed
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
  l2_weight_decay: float = 0.0
  use_groundtruth_dimension: bool = True


@dataclasses.dataclass
class BASNetTask(cfg.TaskConfig):
  """The model config."""
  model: BASNetModel = BASNetModel()
  train_data: DataConfig = DataConfig(is_training=True)
  validation_data: DataConfig = DataConfig(is_training=False)
  losses: Losses = Losses()
  gradient_clip_norm: float = 0.0
  init_checkpoint: Optional[str] = None
  init_checkpoint_modules: Union[
      str, List[str]] = 'backbone'  # all, backbone, and/or decoder


@exp_factory.register_config_factory('basnet')
def basnet() -> cfg.ExperimentConfig:
  """BASNet general."""
  return cfg.ExperimentConfig(
      task=BASNetModel(),
      trainer=cfg.TrainerConfig(),
      restrictions=[
          'task.train_data.is_training != None',
          'task.validation_data.is_training != None'
      ])


# DUTS Dataset
DUTS_TRAIN_EXAMPLES = 10553
DUTS_VAL_EXAMPLES = 5019
89
90
DUTS_INPUT_PATH_BASE_TR = 'DUTS_DATASET'
DUTS_INPUT_PATH_BASE_VAL = 'DUTS_DATASET'
Gunho Park's avatar
Gunho Park committed
91
92
93
94
95


@exp_factory.register_config_factory('basnet_duts')
def basnet_duts() -> cfg.ExperimentConfig:
  """Image segmentation on duts with basnet."""
96
  train_batch_size = 64
Gunho Park's avatar
Gunho Park committed
97
98
99
100
101
102
  eval_batch_size = 16
  steps_per_epoch = DUTS_TRAIN_EXAMPLES // train_batch_size
  config = cfg.ExperimentConfig(
      task=BASNetTask(
          model=BASNetModel(
              input_size=[None, None, 3],
Gunho Park's avatar
Gunho Park committed
103
              use_bias=True,
Gunho Park's avatar
Gunho Park committed
104
105
106
107
108
109
110
              norm_activation=common.NormActivation(
                  activation='relu',
                  norm_momentum=0.99,
                  norm_epsilon=1e-3,
                  use_sync_bn=True)),
          losses=Losses(l2_weight_decay=0),
          train_data=DataConfig(
111
112
113
114
115
              input_path=os.path.join(DUTS_INPUT_PATH_BASE_TR,
                                      'tf_record_train'),
              file_type='tfrecord',
              crop_size=[224, 224],
              output_size=[256, 256],
Gunho Park's avatar
Gunho Park committed
116
117
118
119
              is_training=True,
              global_batch_size=train_batch_size,
          ),
          validation_data=DataConfig(
120
121
122
123
              input_path=os.path.join(DUTS_INPUT_PATH_BASE_VAL,
                                      'tf_record_test'),
              file_type='tfrecord',
              output_size=[256, 256],
Gunho Park's avatar
Gunho Park committed
124
125
126
              is_training=False,
              global_batch_size=eval_batch_size,
          ),
127
          init_checkpoint='gs://cloud-basnet-checkpoints/basnet_encoder_imagenet/ckpt-340306',
Gunho Park's avatar
Gunho Park committed
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
          init_checkpoint_modules='backbone'
      ),
      trainer=cfg.TrainerConfig(
          steps_per_loop=steps_per_epoch,
          summary_interval=steps_per_epoch,
          checkpoint_interval=steps_per_epoch,
          train_steps=300 * steps_per_epoch,
          validation_steps=DUTS_VAL_EXAMPLES // eval_batch_size,
          validation_interval=steps_per_epoch,
          optimizer_config=optimization.OptimizationConfig({
              'optimizer': {
                  'type': 'adam',
                  'adam': {
                      'beta_1': 0.9,
                      'beta_2': 0.999,
                      'epsilon': 1e-8,
                  }
              },
              'learning_rate': {
                  'type': 'constant',
                  'constant': {'learning_rate': 0.001}
              }
          })),
      restrictions=[
          'task.train_data.is_training != None',
          'task.validation_data.is_training != None'
      ])
  return config