(
config_yaml,
displayiters,
saveiters,
maxiters,
max_to_keep=5,
keepdeconvweights=True,
allow_growth=True,
)
| 132 | |
| 133 | |
| 134 | def train( |
| 135 | config_yaml, |
| 136 | displayiters, |
| 137 | saveiters, |
| 138 | maxiters, |
| 139 | max_to_keep=5, |
| 140 | keepdeconvweights=True, |
| 141 | allow_growth=True, |
| 142 | ): |
| 143 | start_path = os.getcwd() |
| 144 | os.chdir(str(Path(config_yaml).parents[0])) # switch to folder of config_yaml (for logging) |
| 145 | setup_logging() |
| 146 | |
| 147 | cfg = load_config(config_yaml) |
| 148 | net_type = cfg["net_type"] |
| 149 | if cfg["dataset_type"] in ("scalecrop", "tensorpack", "deterministic"): |
| 150 | print( |
| 151 | "Switching batchsize to 1, as tensorpack/scalecrop/deterministic loaders " |
| 152 | "do not support batches >1. Use imgaug/default loader." |
| 153 | ) |
| 154 | cfg["batch_size"] = 1 # in case this was edited for analysis.- |
| 155 | |
| 156 | dataset = PoseDatasetFactory.create(cfg) |
| 157 | batch_spec = get_batch_spec(cfg) |
| 158 | batch, enqueue_op, placeholders = setup_preloading(batch_spec) |
| 159 | |
| 160 | losses = PoseNetFactory.create(cfg).train(batch) |
| 161 | total_loss = losses["total_loss"] |
| 162 | |
| 163 | for k, t in losses.items(): |
| 164 | tf.compat.v1.summary.scalar(k, t) |
| 165 | merged_summaries = tf.compat.v1.summary.merge_all() |
| 166 | |
| 167 | stem = Path(cfg["init_weights"]).stem |
| 168 | if "snapshot" in stem and keepdeconvweights: |
| 169 | print("Loading already trained DLC with backbone:", net_type) |
| 170 | variables_to_restore = slim.get_variables_to_restore() |
| 171 | start_iter = int(stem.split("-")[1]) |
| 172 | else: |
| 173 | print("Loading ImageNet-pretrained", net_type) |
| 174 | # loading backbone from ResNet, MobileNet etc. |
| 175 | if "resnet" in net_type: |
| 176 | variables_to_restore = slim.get_variables_to_restore(include=["resnet_v1"]) |
| 177 | elif "mobilenet" in net_type: |
| 178 | variables_to_restore = slim.get_variables_to_restore(include=["MobilenetV2"]) |
| 179 | elif "efficientnet" in net_type: |
| 180 | variables_to_restore = slim.get_variables_to_restore(include=["efficientnet"]) |
| 181 | variables_to_restore = { |
| 182 | var.op.name.replace("efficientnet/", "") + "/ExponentialMovingAverage": var |
| 183 | for var in variables_to_restore |
| 184 | } |
| 185 | else: |
| 186 | print("Wait for DLC 2.3.") |
| 187 | start_iter = 0 |
| 188 | |
| 189 | restorer = tf.compat.v1.train.Saver(variables_to_restore) |
| 190 | saver = tf.compat.v1.train.Saver( |
| 191 | max_to_keep=max_to_keep |
no test coverage detected