MCPcopy Create free account
hub / github.com/Topdu/OpenOCR / _init_rec_model

Method _init_rec_model

tools/engine/trainer.py:164–200  ·  view source on GitHub ↗
(self)

Source from the content-addressed store, hash-verified

162 f'run with torch {torch.__version__} and device {self.device}')
163
164 def _init_rec_model(self):
165 from openrec.losses import build_loss as build_rec_loss
166 from openrec.metrics import build_metric as build_rec_metric
167 from openrec.modeling import build_model as build_rec_model
168 from openrec.postprocess import build_post_process as build_rec_post_process
169
170 # build post process
171 self.post_process_class = build_rec_post_process(
172 self.cfg['PostProcess'], self.cfg['Global'])
173 # build model
174 # for rec algorithm
175 self.use_transformers = self.cfg['Global'].get('use_transformers',
176 False)
177 if self.use_transformers:
178 if self.cfg['Architecture']['algorithm'] == 'UniRec':
179 from openrec.modeling.unirec_modeling.modeling_unirec import UniRecForConditionalGenerationNew
180 from openrec.modeling.unirec_modeling.configuration_unirec import UniRecConfig
181 cfg_vlm = UniRecConfig.from_pretrained(
182 self.cfg['Global']['vlm_ocr_config'])
183 cfg_vlm._attn_implementation = 'flash_attention_2'
184 # cfg_vlm._attn_implementation = "eager"
185 # cfg_vlm._attn_implementation = "sdpa"
186 self.model = UniRecForConditionalGenerationNew(config=cfg_vlm)
187 elif self.cfg['Architecture']['algorithm'] == 'CMER':
188 from openrec.modeling.cmer_modeling.modeling_cmer import CMER, CMERConfig
189 cfg_model = CMERConfig(
190 self.cfg['Architecture']['vision_config'],
191 self.cfg['Architecture']['decoder_config'])
192 self.model = CMER(config=cfg_model)
193 else:
194 char_num = self.post_process_class.get_character_num()
195 self.cfg['Architecture']['Decoder']['out_channels'] = char_num
196 self.model = build_rec_model(self.cfg['Architecture'])
197 # build loss
198 self.loss_class = build_rec_loss(self.cfg['Loss'])
199 # build metric
200 self.eval_class = build_rec_metric(self.cfg['Metric'])
201
202 def _init_det_model(self):
203 from opendet.losses import build_loss as build_det_loss

Callers 1

__init__Method · 0.95

Calls 5

CMERConfigClass · 0.90
CMERClass · 0.90
getMethod · 0.80
get_character_numMethod · 0.45

Tested by

no test coverage detected