Function train

code/ASR/Adapter/train.py:272–315

(dataloaders, model, optimizer, save_path)

    return dict_average(stats)

def train(dataloaders, model, optimizer, save_path):
    train_loader, val_loader, test_loader = dataloaders
    best_loss = float("inf")
    early_stop = 0
    log_json = []
    for epoch in range(args.start_epoch, args.epochs + 1):
        early_stop += 1
        epoch_stats = collections.OrderedDict(epoch=epoch)
        train_stats = train_epoch(train_loader, model, optimizer, epoch)
        valid_stats = test(f"val_{epoch}", val_loader, model, visualize_sim_adapter=args.sim_adapter)
        if best_loss > valid_stats["loss"]:  # Save loss best model
            best_loss = valid_stats["loss"]
            torch_save(model, save_path)
            early_stop = 0

        test_stats = test(f"test_{epoch}", test_loader, model)
        logging.warning(
            f"Epoch: {epoch}, Iteration: {epoch * len(train_loader)}, "
            + f"train loss: {train_stats['loss']:.4f}, dev loss: {valid_stats['loss']:.3f}, test loss: {test_stats['loss']:.3f}"
        )

        torch_save(model, f"{args.outdir}/snapshot.ep.{epoch}", optimizer=optimizer)
        for key in sorted(list(set(list(train_stats.keys()) + list(test_stats.keys())))):
            if not key.endswith("_lst"):
                if key in train_stats:
                    epoch_stats[f"main/{key}"] = train_stats[key]
                if key in valid_stats:
                    epoch_stats[f"validation/main/{key}"] = valid_stats[key]
                if key in test_stats:
                    epoch_stats[f"test/main/{key}"] = test_stats[key]
                
        log_json.append(epoch_stats)
        with open(f"{args.outdir}/log", "w") as f:
            json.dump(log_json, f,
                indent=4,
                ensure_ascii=False,
                separators=(",", ": "),
            )
            logging.warning(f"Log saved at {args.outdir}/log")
            
        if args.patience > 0 and early_stop >= args.patience:
            test_stats = test("test_best", test_loader, model, save_path)
            logging.warning(f"=====Early stop! Final best test loss: {test_stats['loss']}")
            break

if __name__ == "__main__":
    # 执行该命令运行4 GPU训练：CUDA_VISIBLE_DEVICES=0,1,2,3 python -m torch.distributed.launch --nproc_per_node=2 train.py