def test_models_onnx_export(arch_name, input_shape, output_size):
# Model
batch_size = 2
tf.keras.backend.clear_session()
if arch_name == "mobilenet_v3_small_orientation":
model = classification.__dict__[arch_name](pretrained=True, input_shape=input_shape)
else:
model = classification.__dict__[arch_name](pretrained=True, include_top=True, input_shape=input_shape)
# batch_size = None for dynamic batch size
dummy_input = [tf.TensorSpec([None, *input_shape], tf.float32, name="input")]
np_dummy_input = np.random.rand(batch_size, *input_shape).astype(np.float32)
with tempfile.TemporaryDirectory() as tmpdir:
# Export
model_path, output = export_model_to_onnx(
model, model_name=os.path.join(tmpdir, "model"), dummy_input=dummy_input
)
assert os.path.exists(model_path)
# Inference
ort_session = onnxruntime.InferenceSession(
os.path.join(tmpdir, "model.onnx"), providers=["CPUExecutionProvider"]
)
ort_outs = ort_session.run(output, {"input": np_dummy_input})
assert isinstance(ort_outs, list) and len(ort_outs) == 1
assert ort_outs[0].shape == (batch_size, *output_size)
def test_models_onnx_export(arch_name, input_shape, output_size):
# Model
batch_size = 2
model = detection.__dict__[arch_name](pretrained=True, exportable=True).eval()
dummy_input = torch.rand((batch_size, *input_shape), dtype=torch.float32)
with tempfile.TemporaryDirectory() as tmpdir:
# Export
model_path = export_model_to_onnx(model, model_name=os.path.join(tmpdir, "model"), dummy_input=dummy_input)
assert os.path.exists(model_path)
# Inference
ort_session = onnxruntime.InferenceSession(
os.path.join(tmpdir, "model.onnx"), providers=["CPUExecutionProvider"]
)
ort_outs = ort_session.run(["logits"], {"input": dummy_input.numpy()})
assert isinstance(ort_outs, list) and len(ort_outs) == 1
assert ort_outs[0].shape == (batch_size, *output_size)
def test_models_onnx_export(arch_name, input_shape):
# Model
batch_size = 2
tf.keras.backend.clear_session()
model = recognition.__dict__[arch_name](pretrained=True,
exportable=True,
input_shape=input_shape)
if arch_name == "sar_resnet31": # SAR export currently only available with constant batch size
dummy_input = [
tf.TensorSpec([batch_size, *input_shape], tf.float32, name="input")
]
elif arch_name == "master": # MASTER export currently only available with constant batch size
dummy_input = [
tf.TensorSpec([batch_size, *input_shape], tf.float32, name="input")
]
else:
# batch_size = None for dynamic batch size
dummy_input = [
tf.TensorSpec([None, *input_shape], tf.float32, name="input")
]
np_dummy_input = np.random.rand(batch_size,
*input_shape).astype(np.float32)
with tempfile.TemporaryDirectory() as tmpdir:
# Export
model_path, output = export_model_to_onnx(
model,
model_name=os.path.join(tmpdir, "model"),
dummy_input=dummy_input,
large_model=True if arch_name == "master" else False,
)
assert os.path.exists(model_path)
if arch_name == "master":
# large models are exported as zip archive
shutil.unpack_archive(model_path, tmpdir, "zip")
model_path = os.path.join(tmpdir, "__MODEL_PROTO.onnx")
else:
model_path = os.path.join(tmpdir, "model.onnx")
# Inference
ort_session = onnxruntime.InferenceSession(
model_path, providers=["CPUExecutionProvider"])
ort_outs = ort_session.run(output, {"input": np_dummy_input})
assert isinstance(ort_outs, list) and len(ort_outs) == 1
assert ort_outs[0].shape[0] == batch_size
def main(args):
print(args)
if args.push_to_hub:
login_to_hub()
if not isinstance(args.workers, int):
args.workers = min(16, mp.cpu_count())
torch.backends.cudnn.benchmark = True
vocab = VOCABS[args.vocab]
fonts = args.font.split(",")
# Load val data generator
st = time.time()
val_set = CharacterGenerator(
vocab=vocab,
num_samples=args.val_samples * len(vocab),
cache_samples=True,
img_transforms=Compose(
[
T.Resize((args.input_size, args.input_size)),
# Ensure we have a 90% split of white-background images
T.RandomApply(T.ColorInversion(), 0.9),
]
),
font_family=fonts,
)
val_loader = DataLoader(
val_set,
batch_size=args.batch_size,
drop_last=False,
num_workers=args.workers,
sampler=SequentialSampler(val_set),
pin_memory=torch.cuda.is_available(),
)
print(f"Validation set loaded in {time.time() - st:.4}s ({len(val_set)} samples in " f"{len(val_loader)} batches)")
batch_transforms = Normalize(mean=(0.694, 0.695, 0.693), std=(0.299, 0.296, 0.301))
# Load doctr model
model = classification.__dict__[args.arch](pretrained=args.pretrained, num_classes=len(vocab), classes=list(vocab))
# Resume weights
if isinstance(args.resume, str):
print(f"Resuming {args.resume}")
checkpoint = torch.load(args.resume, map_location="cpu")
model.load_state_dict(checkpoint)
# GPU
if isinstance(args.device, int):
if not torch.cuda.is_available():
raise AssertionError("PyTorch cannot access your GPU. Please investigate!")
if args.device >= torch.cuda.device_count():
raise ValueError("Invalid device index")
# Silent default switch to GPU if available
elif torch.cuda.is_available():
args.device = 0
else:
logging.warning("No accessible GPU, targe device set to CPU.")
if torch.cuda.is_available():
torch.cuda.set_device(args.device)
model = model.cuda()
if args.test_only:
print("Running evaluation")
val_loss, acc = evaluate(model, val_loader, batch_transforms)
print(f"Validation loss: {val_loss:.6} (Acc: {acc:.2%})")
return
st = time.time()
# Load train data generator
train_set = CharacterGenerator(
vocab=vocab,
num_samples=args.train_samples * len(vocab),
cache_samples=True,
img_transforms=Compose(
[
T.Resize((args.input_size, args.input_size)),
# Augmentations
T.RandomApply(T.ColorInversion(), 0.9),
# GaussianNoise
T.RandomApply(Grayscale(3), 0.1),
ColorJitter(brightness=0.3, contrast=0.3, saturation=0.3, hue=0.02),
T.RandomApply(GaussianBlur(kernel_size=(3, 3), sigma=(0.1, 3)), 0.3),
RandomRotation(15, interpolation=InterpolationMode.BILINEAR),
]
),
font_family=fonts,
)
train_loader = DataLoader(
train_set,
batch_size=args.batch_size,
drop_last=True,
num_workers=args.workers,
sampler=RandomSampler(train_set),
pin_memory=torch.cuda.is_available(),
)
print(f"Train set loaded in {time.time() - st:.4}s ({len(train_set)} samples in " f"{len(train_loader)} batches)")
if args.show_samples:
x, target = next(iter(train_loader))
plot_samples(x, list(map(vocab.__getitem__, target)))
return
# Optimizer
optimizer = torch.optim.Adam(
[p for p in model.parameters() if p.requires_grad],
args.lr,
betas=(0.95, 0.99),
eps=1e-6,
weight_decay=args.weight_decay,
)
# LR Finder
if args.find_lr:
lrs, losses = record_lr(model, train_loader, batch_transforms, optimizer, amp=args.amp)
plot_recorder(lrs, losses)
return
# Scheduler
if args.sched == "cosine":
scheduler = CosineAnnealingLR(optimizer, args.epochs * len(train_loader), eta_min=args.lr / 25e4)
elif args.sched == "onecycle":
scheduler = OneCycleLR(optimizer, args.lr, args.epochs * len(train_loader))
# Training monitoring
current_time = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
exp_name = f"{args.arch}_{current_time}" if args.name is None else args.name
# W&B
if args.wb:
run = wandb.init(
name=exp_name,
project="character-classification",
config={
"learning_rate": args.lr,
"epochs": args.epochs,
"weight_decay": args.weight_decay,
"batch_size": args.batch_size,
"architecture": args.arch,
"input_size": args.input_size,
"optimizer": "adam",
"framework": "pytorch",
"vocab": args.vocab,
"scheduler": args.sched,
"pretrained": args.pretrained,
},
)
# Create loss queue
min_loss = np.inf
# Training loop
mb = master_bar(range(args.epochs))
for epoch in mb:
fit_one_epoch(model, train_loader, batch_transforms, optimizer, scheduler, mb)
# Validation loop at the end of each epoch
val_loss, acc = evaluate(model, val_loader, batch_transforms)
if val_loss < min_loss:
print(f"Validation loss decreased {min_loss:.6} --> {val_loss:.6}: saving state...")
torch.save(model.state_dict(), f"./{exp_name}.pt")
min_loss = val_loss
mb.write(f"Epoch {epoch + 1}/{args.epochs} - Validation loss: {val_loss:.6} (Acc: {acc:.2%})")
# W&B
if args.wb:
wandb.log(
{
"val_loss": val_loss,
"acc": acc,
}
)
if args.wb:
run.finish()
if args.push_to_hub:
push_to_hf_hub(model, exp_name, task="classification", run_config=args)
if args.export_onnx:
print("Exporting model to ONNX...")
dummy_batch = next(iter(val_loader))
dummy_input = dummy_batch[0].cuda() if torch.cuda.is_available() else dummy_batch[0]
model_path = export_model_to_onnx(model, exp_name, dummy_input)
print(f"Exported model saved in {model_path}")
def main(args):
print(args)
if args.push_to_hub:
login_to_hub()
if not isinstance(args.workers, int):
args.workers = min(16, mp.cpu_count())
vocab = VOCABS[args.vocab]
fonts = args.font.split(",")
# AMP
if args.amp:
mixed_precision.set_global_policy("mixed_float16")
# Load val data generator
st = time.time()
val_set = CharacterGenerator(
vocab=vocab,
num_samples=args.val_samples * len(vocab),
cache_samples=True,
img_transforms=T.Compose(
[
T.Resize((args.input_size, args.input_size)),
# Ensure we have a 90% split of white-background images
T.RandomApply(T.ColorInversion(), 0.9),
]
),
font_family=fonts,
)
val_loader = DataLoader(
val_set,
batch_size=args.batch_size,
shuffle=False,
drop_last=False,
num_workers=args.workers,
collate_fn=collate_fn,
)
print(
f"Validation set loaded in {time.time() - st:.4}s ({len(val_set)} samples in "
f"{val_loader.num_batches} batches)"
)
# Load doctr model
model = classification.__dict__[args.arch](
pretrained=args.pretrained,
input_shape=(args.input_size, args.input_size, 3),
num_classes=len(vocab),
classes=list(vocab),
include_top=True,
)
# Resume weights
if isinstance(args.resume, str):
model.load_weights(args.resume)
batch_transforms = T.Compose(
[
T.Normalize(mean=(0.694, 0.695, 0.693), std=(0.299, 0.296, 0.301)),
]
)
if args.test_only:
print("Running evaluation")
val_loss, acc = evaluate(model, val_loader, batch_transforms)
print(f"Validation loss: {val_loss:.6} (Acc: {acc:.2%})")
return
st = time.time()
# Load train data generator
train_set = CharacterGenerator(
vocab=vocab,
num_samples=args.train_samples * len(vocab),
cache_samples=True,
img_transforms=T.Compose(
[
T.Resize((args.input_size, args.input_size)),
# Augmentations
T.RandomApply(T.ColorInversion(), 0.9),
T.RandomApply(T.ToGray(3), 0.1),
T.RandomJpegQuality(60),
T.RandomSaturation(0.3),
T.RandomContrast(0.3),
T.RandomBrightness(0.3),
# Blur
T.RandomApply(T.GaussianBlur(kernel_shape=(3, 3), std=(0.1, 3)), 0.3),
]
),
font_family=fonts,
)
train_loader = DataLoader(
train_set,
batch_size=args.batch_size,
shuffle=True,
drop_last=True,
num_workers=args.workers,
collate_fn=collate_fn,
)
print(
f"Train set loaded in {time.time() - st:.4}s ({len(train_set)} samples in "
f"{train_loader.num_batches} batches)"
)
if args.show_samples:
x, target = next(iter(train_loader))
plot_samples(x, list(map(vocab.__getitem__, target)))
return
# Optimizer
scheduler = tf.keras.optimizers.schedules.ExponentialDecay(
args.lr,
decay_steps=args.epochs * len(train_loader),
decay_rate=1 / (1e3), # final lr as a fraction of initial lr
staircase=False,
)
optimizer = tf.keras.optimizers.Adam(
learning_rate=scheduler,
beta_1=0.95,
beta_2=0.99,
epsilon=1e-6,
)
if args.amp:
optimizer = mixed_precision.LossScaleOptimizer(optimizer)
# LR Finder
if args.find_lr:
lrs, losses = record_lr(model, train_loader, batch_transforms, optimizer, amp=args.amp)
plot_recorder(lrs, losses)
return
# Tensorboard to monitor training
current_time = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
exp_name = f"{args.arch}_{current_time}" if args.name is None else args.name
# W&B
if args.wb:
run = wandb.init(
name=exp_name,
project="character-classification",
config={
"learning_rate": args.lr,
"epochs": args.epochs,
"weight_decay": 0.0,
"batch_size": args.batch_size,
"architecture": args.arch,
"input_size": args.input_size,
"optimizer": "adam",
"framework": "tensorflow",
"vocab": args.vocab,
"scheduler": "exp_decay",
"pretrained": args.pretrained,
},
)
# Create loss queue
min_loss = np.inf
# Training loop
mb = master_bar(range(args.epochs))
for epoch in mb:
fit_one_epoch(model, train_loader, batch_transforms, optimizer, mb, args.amp)
# Validation loop at the end of each epoch
val_loss, acc = evaluate(model, val_loader, batch_transforms)
if val_loss < min_loss:
print(f"Validation loss decreased {min_loss:.6} --> {val_loss:.6}: saving state...")
model.save_weights(f"./{exp_name}/weights")
min_loss = val_loss
mb.write(f"Epoch {epoch + 1}/{args.epochs} - Validation loss: {val_loss:.6} (Acc: {acc:.2%})")
# W&B
if args.wb:
wandb.log(
{
"val_loss": val_loss,
"acc": acc,
}
)
if args.wb:
run.finish()
if args.push_to_hub:
push_to_hf_hub(model, exp_name, task="classification", run_config=args)
if args.export_onnx:
print("Exporting model to ONNX...")
dummy_input = [tf.TensorSpec([None, args.input_size, args.input_size, 3], tf.float32, name="input")]
model_path, _ = export_model_to_onnx(model, exp_name, dummy_input)
print(f"Exported model saved in {model_path}")