Options -> ModelConfig

This commit is contained in:
jamjamjon
2025-05-16 15:45:27 +08:00
parent 54045543fc
commit 70e6b2f03e
118 changed files with 1735 additions and 2503 deletions

View File

@ -1,6 +1,6 @@
use usls::{
models::{TrOCR, TrOCRKind},
DataLoader, Options, Scale,
DataLoader, ModelConfig, Scale,
};
#[derive(argh::FromArgs)]
@ -38,52 +38,22 @@ fn main() -> anyhow::Result<()> {
])?;
// build model
let (options_encoder, options_decoder, options_decoder_merged) =
match args.scale.as_str().try_into()? {
Scale::S => match args.kind.as_str().try_into()? {
TrOCRKind::Printed => (
Options::trocr_encoder_small_printed(),
Options::trocr_decoder_small_printed(),
Options::trocr_decoder_merged_small_printed(),
),
TrOCRKind::HandWritten => (
Options::trocr_encoder_small_handwritten(),
Options::trocr_decoder_small_handwritten(),
Options::trocr_decoder_merged_small_handwritten(),
),
},
Scale::B => match args.kind.as_str().try_into()? {
TrOCRKind::Printed => (
Options::trocr_encoder_base_printed(),
Options::trocr_decoder_base_printed(),
Options::trocr_decoder_merged_base_printed(),
),
TrOCRKind::HandWritten => (
Options::trocr_encoder_base_handwritten(),
Options::trocr_decoder_base_handwritten(),
Options::trocr_decoder_merged_base_handwritten(),
),
},
x => anyhow::bail!("Unsupported TrOCR scale: {:?}", x),
};
let config = match args.scale.as_str().try_into()? {
Scale::S => match args.kind.as_str().try_into()? {
TrOCRKind::Printed => ModelConfig::trocr_small_printed(),
TrOCRKind::HandWritten => ModelConfig::trocr_small_handwritten(),
},
Scale::B => match args.kind.as_str().try_into()? {
TrOCRKind::Printed => ModelConfig::trocr_base_printed(),
TrOCRKind::HandWritten => ModelConfig::trocr_base_handwritten(),
},
x => anyhow::bail!("Unsupported TrOCR scale: {:?}", x),
}
.with_device_all(args.device.as_str().try_into()?)
.with_dtype_all(args.dtype.as_str().try_into()?)
.commit()?;
let mut model = TrOCR::new(
options_encoder
.with_model_device(args.device.as_str().try_into()?)
.with_model_dtype(args.dtype.as_str().try_into()?)
.with_batch_size(xs.len())
.commit()?,
options_decoder
.with_model_device(args.device.as_str().try_into()?)
.with_model_dtype(args.dtype.as_str().try_into()?)
.with_batch_size(xs.len())
.commit()?,
options_decoder_merged
.with_model_device(args.device.as_str().try_into()?)
.with_model_dtype(args.dtype.as_str().try_into()?)
.with_batch_size(xs.len())
.commit()?,
)?;
let mut model = TrOCR::new(config)?;
// inference
let ys = model.forward(&xs)?;