helper function to create model from given configurations
(opt)
| 58 | |
| 59 | |
| 60 | def get_model(opt): |
| 61 | """ |
| 62 | helper function to create model from given configurations |
| 63 | """ |
| 64 | if opt["model_type"] == "SegmentationModel": |
| 65 | seg_model = PCBSegNet(opt) |
| 66 | model = seg_model.build() |
| 67 | optimizer = getattr(tf.keras.optimizers, opt["train"]["optim"]["type"])( |
| 68 | learning_rate=opt["train"]["optim"]["lr"], |
| 69 | beta_1=opt["train"]["optim"]["betas"][0], |
| 70 | beta_2=opt["train"]["optim"]["betas"][1], |
| 71 | ) |
| 72 | loss = getattr(models, opt["train"]["loss"]["type"]) |
| 73 | metrics = [ |
| 74 | getattr(models, opt["train"]["metric"][item]["type"]) |
| 75 | for item in opt["train"]["metric"] |
| 76 | ] |
| 77 | model.compile(optimizer=optimizer, loss=loss, metrics=metrics) |
| 78 | return model |
| 79 | |
| 80 | elif opt["model_type"] == "ClassificationModel": |
| 81 | class_model = PCBClassNet(opt) |
| 82 | model = class_model.build() |
| 83 | optimizer = getattr(tf.keras.optimizers, opt["train"]["optim"]["type"])( |
| 84 | learning_rate=opt["train"]["optim"]["lr"], |
| 85 | beta_1=opt["train"]["optim"]["betas"][0], |
| 86 | beta_2=opt["train"]["optim"]["betas"][1], |
| 87 | ) |
| 88 | # tf.keras.metrics.Precision() |
| 89 | metrics = [ |
| 90 | getattr(tf.keras.metrics, opt["train"]["metric"][item]["type"])() |
| 91 | for item in opt["train"]["metric"] |
| 92 | ] + ["accuracy"] |
| 93 | model.compile( |
| 94 | optimizer=optimizer, loss=opt["train"]["loss"]["type"], metrics=metrics |
| 95 | ) |
| 96 | return model |
| 97 | |
| 98 | else: |
| 99 | assert ( |
| 100 | False |
| 101 | ), f"Found model type as {opt['model_type']} \ |
| 102 | but it should be one of SegmentationModel/ClassificationModel" |