MCPcopy Create free account
hub / github.com/CandleLabAI/PCBSegClassNet / main

Function main

src/train_segmentation.py:60–128  ·  view source on GitHub ↗

main function for segmentation training

()

Source from the content-addressed store, hash-verified

58 logger.info(dict2str(opt))
59
60def main():
61 """
62 main function for segmentation training
63 """
64 opt = parse_config()
65 init_log(opt)
66
67 # get dataloaders
68 train_dataset, val_dataset = get_data(opt)
69 logger = logging.getLogger(__name__)
70 logger.info(f"Found {len(train_dataset)} batches for training")
71 logger.info(f"Found {len(val_dataset)} batches for validation")
72
73 # get model
74 model = get_model(opt)
75 logger.info(model.summary())
76
77 callbacks = []
78 if "modelcheckpoint" in opt["train"]["callbacks"]:
79 logger.info("Using modelcheckpoint callback")
80 callbacks.append(
81 getattr(
82 tf.keras.callbacks, opt["train"]["callbacks"]["modelcheckpoint"]["type"]
83 )(
84 filepath=opt["path"]["checkpoint_network"],
85 monitor=opt["train"]["callbacks"]["modelcheckpoint"]["monitor"],
86 mode=opt["train"]["callbacks"]["modelcheckpoint"]["mode"],
87 verbose=opt["train"]["callbacks"]["modelcheckpoint"]["verbose"],
88 save_best_only=opt["train"]["callbacks"]["modelcheckpoint"][
89 "save_best_only"
90 ],
91 save_weights_only=opt["train"]["callbacks"]["modelcheckpoint"][
92 "save_weights_only"
93 ],
94 )
95 )
96
97 if "reducelronplateau" in opt["train"]["callbacks"]:
98 logger.info("Using reducelronplateau callback")
99 callbacks.append(
100 getattr(
101 tf.keras.callbacks,
102 opt["train"]["callbacks"]["reducelronplateau"]["type"],
103 )(
104 monitor=opt["train"]["callbacks"]["reducelronplateau"]["monitor"],
105 mode=opt["train"]["callbacks"]["reducelronplateau"]["mode"],
106 verbose=opt["train"]["callbacks"]["reducelronplateau"]["verbose"],
107 factor=opt["train"]["callbacks"]["reducelronplateau"]["factor"],
108 patience=opt["train"]["callbacks"]["reducelronplateau"]["patience"],
109 min_lr=opt["train"]["callbacks"]["reducelronplateau"]["min_lr"],
110 )
111 )
112
113 # training model
114 if opt["train"]["total_epochs"] > 0:
115 logger.info(f"Training for {opt['train']['total_epochs']} epochs")
116 model.fit(
117 train_dataset,

Callers 1

Calls 4

get_dataFunction · 0.90
get_modelFunction · 0.90
parse_configFunction · 0.70
init_logFunction · 0.70

Tested by

no test coverage detected