import torch import torch.utils.data import torch.backends.cudnn as cudnn from torch.utils.tensorboard import SummaryWriter import os import json import valid from utils import utils from utils import sam from utils import option from data import dataset from model import htr_convtext from functools import partial import random import numpy as np import re import importlib from model.tcm_head import TCMHead, build_tcm_vocab, make_context_batch import wandb def compute_losses( args, model, tcm_head, image, texts, batch_size, criterion_ctc, converter, nb_iter, ctc_lambda, tcm_lambda, stoi, mask_mode='span', mask_ratio=0.30, block_span=4, max_span_length=8, pre_tcm_ctx=None, use_masking=True, ): if tcm_head is None or nb_iter < args.tcm_warmup_iters: preds = model(image, use_masking=use_masking, mask_mode=mask_mode, mask_ratio=mask_ratio, max_span_length=max_span_length) feats = None else: preds, feats, vis_mask = model( image, use_masking=use_masking, return_features=True, return_mask=True, mask_mode=mask_mode, mask_ratio=mask_ratio, block_span=block_span, max_span_length=max_span_length ) text_ctc, length_ctc = converter.encode(texts) text_ctc = text_ctc.to(preds.device) length_ctc = length_ctc.to(preds.device) preds_sz = torch.full((batch_size,), preds.size( 1), dtype=torch.int32, device=preds.device) loss_ctc = criterion_ctc(preds.permute(1, 0, 2).log_softmax(2), text_ctc, preds_sz, length_ctc).mean() loss_tcm = torch.zeros((), device=preds.device) if tcm_head is not None and feats is not None: left_ctx, right_ctx, tgt_ids, tgt_mask = pre_tcm_ctx if pre_tcm_ctx is not None else make_context_batch( texts, stoi, sub_str_len=args.tcm_sub_len, device=image.device) if vis_mask is not None: B_v, N_v = vis_mask.shape B_t, L_t = tgt_mask.shape if N_v != L_t: idx = torch.linspace(0, N_v - 1, steps=L_t, device=vis_mask.device).long() focus_mask = vis_mask[:, idx] else: focus_mask = vis_mask else: focus_mask = None out = tcm_head( feats, left_ctx, right_ctx, tgt_ids, tgt_mask, focus_mask=focus_mask ) loss_tcm = out['loss_tcm'] total = ctc_lambda * loss_ctc + tcm_lambda * loss_tcm return total, loss_ctc.detach(), loss_tcm.detach() def tri_masked_loss(args, model, tcm_head, image, labels, batch_size, criterion, converter, nb_iter, ctc_lambda, tcm_lambda, stoi, r_rand=0.6, r_block=0.6, block_span=4, r_span=0.4, max_span=8): total = 0.0 total_ctc = 0.0 total_tcm = 0.0 plans = [("random", r_rand), ("block", r_block), ("span", r_span)] if tcm_head is not None and nb_iter >= args.tcm_warmup_iters: pre_tcm_ctx = make_context_batch( labels, stoi, sub_str_len=args.tcm_sub_len, device=image.device) for mode, ratio in plans: loss, loss_ctc, loss_tcm = compute_losses( args, model, tcm_head, image, labels, batch_size, criterion, converter, nb_iter, ctc_lambda, tcm_lambda, stoi, mask_mode=mode, mask_ratio=ratio, block_span=block_span, max_span_length=max_span, pre_tcm_ctx=pre_tcm_ctx ) total += loss total_ctc += loss_ctc total_tcm += loss_tcm denom = 3.0 return total/denom, total_ctc/denom, total_tcm/denom def main(): args = option.get_args_parser() torch.manual_seed(args.seed) args.save_dir = os.path.join(args.out_dir, args.exp_name) os.makedirs(args.save_dir, exist_ok=True) logger = utils.get_logger(args.save_dir) logger.info(json.dumps(vars(args), indent=4, sort_keys=True)) writer = SummaryWriter(args.save_dir) if getattr(args, 'use_wandb', False): try: wandb = importlib.import_module('wandb') wandb.init(project=getattr(args, 'wandb_project', 'None'), name=args.exp_name, config=vars(args), dir=args.save_dir) logger.info("Weights & Biases logging enabled") except Exception as e: logger.warning( f"Failed to initialize wandb: {e}. Continuing without wandb.") wandb = None else: wandb = None torch.backends.cudnn.benchmark = True model = htr_convtext.create_model( nb_cls=args.nb_cls, img_size=args.img_size[::-1]) total_param = sum(p.numel() for p in model.parameters()) logger.info('total_param is {}'.format(total_param)) model.train() model = model.cuda() ema_decay = args.ema_decay logger.info(f"Using EMA decay: {ema_decay}") model_ema = utils.ModelEma(model, ema_decay) model.zero_grad() resume_path = args.resume best_cer, best_wer, start_iter, optimizer_state, train_loss, train_loss_count = utils.load_checkpoint( model, model_ema, None, resume_path, logger) logger.info('Loading train loader...') train_dataset = dataset.myLoadDS( args.train_data_list, args.data_path, args.img_size, dataset=args.dataset) train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=args.train_bs, shuffle=True, pin_memory=True, num_workers=args.num_workers, collate_fn=partial(dataset.SameTrCollate, args=args)) train_iter = dataset.cycle_data(train_loader) logger.info('Loading val loader...') val_dataset = dataset.myLoadDS( args.val_data_list, args.data_path, args.img_size, ralph=train_dataset.ralph, dataset=args.dataset) val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=args.val_bs, shuffle=False, pin_memory=True, num_workers=args.num_workers) criterion = torch.nn.CTCLoss(reduction='none', zero_infinity=True) converter = utils.CTCLabelConverter(train_dataset.ralph.values()) stoi, itos, pad_id = build_tcm_vocab(converter) vocab_size_tcm = len(itos) d_vis = model.embed_dim if args.tcm_enable: tcm_head = TCMHead(d_vis=d_vis, vocab_size_tcm=vocab_size_tcm, pad_id=pad_id, sub_str_len=args.tcm_sub_len).cuda() tcm_head.train() else: tcm_head = None param_groups = list(model.parameters()) if args.tcm_enable and tcm_head is not None: param_groups += list(tcm_head.parameters()) logger.info( f"Optimizing {sum(p.numel() for p in tcm_head.parameters())} tcm params in addition to model params") optimizer = sam.SAM(param_groups, torch.optim.AdamW, lr=1e-7, betas=(0.9, 0.99), weight_decay=args.weight_decay) if optimizer_state is not None: try: optimizer.load_state_dict(optimizer_state) logger.info("Successfully loaded optimizer state") except Exception as e: logger.warning(f"Failed to load optimizer state: {e}") logger.info( "Continuing training without optimizer state (will restart from initial lr/momentum)") elif resume_path and os.path.isfile(resume_path): try: ckpt = torch.load(resume_path, map_location='cpu', weights_only=False) if 'optimizer' in ckpt: optimizer.load_state_dict(ckpt['optimizer']) logger.info("Loaded optimizer state from checkpoint directly") except Exception as e: logger.warning( f"Could not load optimizer state from checkpoint: {e}") if resume_path and os.path.isfile(resume_path) and tcm_head is not None: try: ckpt = torch.load(resume_path, map_location='cpu', weights_only=False) if 'tcm_head' in ckpt: tcm_head.load_state_dict(ckpt['tcm_head'], strict=False) logger.info("Restored tcm head state from checkpoint") else: logger.info( "No tcm head state found in checkpoint; training tcm from scratch") except Exception as e: logger.warning(f"Failed to restore tcm head from checkpoint: {e}") best_cer, best_wer = best_cer, best_wer train_loss = train_loss train_loss_count = train_loss_count #### ---- train & eval ---- #### logger.info('Start training...') accum_steps = max(1, int(getattr(args, 'accum_steps', 1))) micro_step = 0 avg_loss_ctc = 0.0 avg_loss_tcm = 0.0 for nb_iter in range(start_iter, args.total_iter): optimizer, current_lr = utils.update_lr_cos( nb_iter, args.warm_up_iter, args.total_iter, args.max_lr, optimizer) optimizer.zero_grad() total_loss_this_macro = 0.0 avg_loss_ctc = 0.0 avg_loss_tcm = 0.0 cached_batches = [] for micro_step in range(accum_steps): batch = next(train_iter) cached_batches.append(batch) image = batch[0].cuda(non_blocking=True) text, length = converter.encode(batch[1]) batch_size = image.size(0) if args.use_masking: # loss, loss_ctc, loss_tcm = tri_masked_loss( # args, model, tcm_head, image, batch[1], batch_size, criterion, converter, # nb_iter, args.ctc_lambda, args.tcm_lambda, stoi, # r_rand=args.r_rand, # r_block=args.r_block, # block_span=args.block_span, # r_span=args.r_span, # max_span=args.max_span # ) loss, loss_ctc, loss_tcm = compute_losses( args, model, tcm_head, image, batch[1], batch_size, criterion, converter, nb_iter, args.ctc_lambda, args.tcm_lambda, stoi, mask_mode='span', mask_ratio=0.4, max_span_length=8, use_masking=True ) else: loss, loss_ctc, loss_tcm = compute_losses( args, model, tcm_head, image, batch[1], batch_size, criterion, converter, nb_iter, args.ctc_lambda, args.tcm_lambda, stoi, use_masking=False ) (loss / accum_steps).backward() total_loss_this_macro += loss.item() avg_loss_ctc += loss_ctc.mean().item() avg_loss_tcm += loss_tcm.mean().item() optimizer.first_step(zero_grad=True) # Recompute with perturbed weights and accumulate again for the second step for micro_step in range(accum_steps): batch = cached_batches[micro_step] image = batch[0].cuda(non_blocking=True) text, length = converter.encode(batch[1]) batch_size = image.size(0) if args.use_masking: # loss2, loss_ctc, loss_tcm = tri_masked_loss( # args, model, tcm_head, image, batch[1], batch_size, criterion, converter, # nb_iter, args.ctc_lambda, args.tcm_lambda, stoi, # r_rand=args.r_rand, # r_block=args.r_block, # block_span=args.block_span, # r_span=args.r_span, # max_span=args.max_span # ) loss2, loss_ctc, loss_tcm = compute_losses( args, model, tcm_head, image, batch[1], batch_size, criterion, converter, nb_iter, args.ctc_lambda, args.tcm_lambda, stoi, mask_mode='span', mask_ratio=0.4, max_span_length=8, use_masking=True ) else: loss2, loss_ctc, loss_tcm = compute_losses( args, model, tcm_head, image, batch[1], batch_size, criterion, converter, nb_iter, args.ctc_lambda, args.tcm_lambda, stoi, use_masking=False ) (loss2 / accum_steps).backward() optimizer.second_step(zero_grad=True) model.zero_grad() model_ema.update(model, num_updates=nb_iter / 2) train_loss += total_loss_this_macro / accum_steps train_loss_count += 1 if nb_iter % args.print_iter == 0: train_loss_avg = train_loss / train_loss_count if train_loss_count > 0 else 0.0 logger.info( f'Iter : {nb_iter} \t LR : {current_lr:0.5f} \t total : {train_loss_avg:0.5f} \t CTC : {(avg_loss_ctc/accum_steps):0.5f} \t tcm : {(avg_loss_tcm/accum_steps):0.5f} \t ') writer.add_scalar('./Train/lr', current_lr, nb_iter) writer.add_scalar('./Train/train_loss', train_loss_avg, nb_iter) if wandb is not None: wandb.log({ 'train/lr': current_lr, 'train/loss': train_loss_avg, 'train/CTC': (avg_loss_ctc/accum_steps), 'train/tcm': (avg_loss_tcm/accum_steps), 'iter': nb_iter, }, step=nb_iter) train_loss = 0.0 train_loss_count = 0 if nb_iter % args.eval_iter == 0: model.eval() with torch.no_grad(): val_loss, val_cer, val_wer, preds, labels = valid.validation(model_ema.ema, criterion, val_loader, converter) if nb_iter % args.eval_iter*5 == 0: ckpt_name = f"checkpoint_{best_cer:.4f}_{best_wer:.4f}_{nb_iter}.pth" checkpoint = { 'model': model.state_dict(), 'state_dict_ema': model_ema.ema.state_dict(), 'optimizer': optimizer.state_dict(), 'nb_iter': nb_iter, 'best_cer': best_cer, 'best_wer': best_wer, 'args': vars(args), 'random_state': random.getstate(), 'numpy_state': np.random.get_state(), 'torch_state': torch.get_rng_state(), 'torch_cuda_state': torch.cuda.get_rng_state() if torch.cuda.is_available() else None, 'train_loss': train_loss, 'train_loss_count': train_loss_count, } if tcm_head is not None: checkpoint['tcm_head'] = tcm_head.state_dict() torch.save(checkpoint, os.path.join( args.save_dir, ckpt_name)) if val_cer < best_cer: logger.info( f'CER improved from {best_cer:.4f} to {val_cer:.4f}!!!') best_cer = val_cer checkpoint = { 'model': model.state_dict(), 'state_dict_ema': model_ema.ema.state_dict(), 'optimizer': optimizer.state_dict(), 'nb_iter': nb_iter, 'best_cer': best_cer, 'best_wer': best_wer, 'args': vars(args), 'random_state': random.getstate(), 'numpy_state': np.random.get_state(), 'torch_state': torch.get_rng_state(), 'torch_cuda_state': torch.cuda.get_rng_state() if torch.cuda.is_available() else None, 'train_loss': train_loss, 'train_loss_count': train_loss_count, } if tcm_head is not None: checkpoint['tcm_head'] = tcm_head.state_dict() torch.save(checkpoint, os.path.join( args.save_dir, 'best_CER.pth')) if val_wer < best_wer: logger.info( f'WER improved from {best_wer:.4f} to {val_wer:.4f}!!!') best_wer = val_wer checkpoint = { 'model': model.state_dict(), 'state_dict_ema': model_ema.ema.state_dict(), 'optimizer': optimizer.state_dict(), 'nb_iter': nb_iter, 'best_cer': best_cer, 'best_wer': best_wer, 'args': vars(args), 'random_state': random.getstate(), 'numpy_state': np.random.get_state(), 'torch_state': torch.get_rng_state(), 'torch_cuda_state': torch.cuda.get_rng_state() if torch.cuda.is_available() else None, 'train_loss': train_loss, 'train_loss_count': train_loss_count, } if tcm_head is not None: checkpoint['tcm_head'] = tcm_head.state_dict() torch.save(checkpoint, os.path.join( args.save_dir, 'best_WER.pth')) logger.info( f'Val. loss : {val_loss:0.3f} \t CER : {val_cer:0.4f} \t WER : {val_wer:0.4f} \t ') writer.add_scalar('./VAL/CER', val_cer, nb_iter) writer.add_scalar('./VAL/WER', val_wer, nb_iter) writer.add_scalar('./VAL/bestCER', best_cer, nb_iter) writer.add_scalar('./VAL/bestWER', best_wer, nb_iter) writer.add_scalar('./VAL/val_loss', val_loss, nb_iter) if wandb is not None: wandb.log({ 'val/loss': val_loss, 'val/CER': val_cer, 'val/WER': val_wer, 'val/best_CER': best_cer, 'val/best_WER': best_wer, 'iter': nb_iter, }, step=nb_iter) model.train() if __name__ == '__main__': main()