|
- import os
- import time
- import argparse
- import random
- import torch
- import numpy as np
- from tensorboardX import SummaryWriter
-
- from utils import load_config, save_checkpoint, load_checkpoint
- from dataset import get_crohme_dataset
- from models.can import CAN
- from training import train, eval
-
- parser = argparse.ArgumentParser(description='model training')
- parser.add_argument('--dataset', default='CROHME', type=str, help='数据集名称')
- parser.add_argument('--check', action='store_true', help='测试代码选项')
- args = parser.parse_args()
-
- if not args.dataset:
- print('请提供数据集名称')
- exit(-1)
-
- if args.dataset == 'CROHME':
- config_file = 'config.yaml'
-
- """加载config文件"""
- params = load_config(config_file)
-
- """设置随机种子"""
- random.seed(params['seed'])
- np.random.seed(params['seed'])
- torch.manual_seed(params['seed'])
- torch.cuda.manual_seed(params['seed'])
-
- os.environ['CUDA_VISIBLE_DEVICES'] = '0'
- device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
- params['device'] = device
-
- if args.dataset == 'CROHME':
- train_loader, eval_loader = get_crohme_dataset(params)
-
- model = CAN(params)
- now = time.strftime("%Y-%m-%d-%H-%M", time.localtime())
- model.name = f'{params["experiment"]}_{now}_decoder-{params["decoder"]["net"]}'
-
- print(model.name)
- model = model.to(device)
-
- if args.check:
- writer = None
- else:
- writer = SummaryWriter(f'{params["log_dir"]}/{model.name}')
-
- optimizer = getattr(torch.optim, params['optimizer'])(model.parameters(), lr=float(params['lr']),
- eps=float(params['eps']), weight_decay=float(params['weight_decay']))
-
- if params['finetune']:
- print('加载预训练模型权重')
- print(f'预训练权重路径: {params["checkpoint"]}')
- load_checkpoint(model, optimizer, params['checkpoint'])
-
- if not args.check:
- if not os.path.exists(os.path.join(params['checkpoint_dir'], model.name)):
- os.makedirs(os.path.join(params['checkpoint_dir'], model.name), exist_ok=True)
- os.system(f'cp {config_file} {os.path.join(params["checkpoint_dir"], model.name, model.name)}.yaml')
-
- """在CROHME上训练"""
- if args.dataset == 'CROHME':
- min_score, init_epoch = 0, 0
-
- for epoch in range(init_epoch, params['epochs']):
- train_loss, train_word_score, train_exprate = train(params, model, optimizer, epoch, train_loader, writer=writer)
-
- if epoch >= params['valid_start']:
- eval_loss, eval_word_score, eval_exprate = eval(params, model, epoch, eval_loader, writer=writer)
- print(f'Epoch: {epoch+1} loss: {eval_loss:.4f} word score: {eval_word_score:.4f} ExpRate: {eval_exprate:.4f}')
- if eval_exprate > min_score and not args.check and epoch >= params['save_start']:
- min_score = eval_exprate
- save_checkpoint(model, optimizer, eval_word_score, eval_exprate, epoch+1,
- optimizer_save=params['optimizer_save'], path=params['checkpoint_dir'])
|