Train the model.
(model)
| 176 | |
| 177 | |
| 178 | def train(model): |
| 179 | """Train the model.""" |
| 180 | # Training dataset. |
| 181 | dataset_train = BalloonDataset() |
| 182 | dataset_train.load_balloon(args.dataset, "train") |
| 183 | dataset_train.prepare() |
| 184 | |
| 185 | # Validation dataset |
| 186 | dataset_val = BalloonDataset() |
| 187 | dataset_val.load_balloon(args.dataset, "val") |
| 188 | dataset_val.prepare() |
| 189 | |
| 190 | # *** This training schedule is an example. Update to your needs *** |
| 191 | # Since we're using a very small dataset, and starting from |
| 192 | # COCO trained weights, we don't need to train too long. Also, |
| 193 | # no need to train all layers, just the heads should do it. |
| 194 | print("Training network heads") |
| 195 | model.train(dataset_train, dataset_val, |
| 196 | learning_rate=config.LEARNING_RATE, |
| 197 | epochs=30, |
| 198 | layers='heads') |
| 199 | |
| 200 | |
| 201 | def color_splash(image, mask): |
no test coverage detected