『ignite』模型的训练过程

最新推荐文章于 2023-09-24 21:31:18 发布

原创

最新推荐文章于 2023-09-24 21:31:18 发布 · 890 阅读

1 ·

CC 4.0 BY-SA版权

文章标签：

#深度学习 #python #ignite #pytorch

本文介绍了一个基于PyTorch的通用训练框架实现，该框架包括模型训练、验证及早停策略等关键组件。通过继承自基类`Trainer`的具体实现`ConvLSTMTrainer`，展示了如何为ConvLSTM模型定制训练流程。

trainer的父类

from typing import Mapping, Dict, Optional

import torch
from ignite.contrib.handlers import ProgressBar
from ignite.engine import Engine, Events
from ignite.handlers import TerminateOnNan, EarlyStopping
from ignite.metrics import Loss, RunningAverage

from src.exception import ModelNotFoundException
from src.experiment import Number


class Trainer(object):

    def __init__(self, *, model: torch.nn.Module = None, file: str = None, save: str = None, device: str = None):
        if model is not None:
            self.model = model
        elif file is not None:
            self.model = torch.load(file, map_location=device)
        else:
            raise ModelNotFoundException("模型未定义，请传入 torch.nn.Module 对象或可加载的模型的文件路径.")

        if device is not None:
            self.device = device
        else:
            self.device = "cuda" if torch.cuda.is_available() else "cpu"

        if save is not None:
            self.save: str = save
        else:
            raise ValueError("模型存储路径未定义！")

        # 评价指标，一般都至少使用 MSE
        self.metrics: Dict = {
   
   "MSE": Loss(torch.nn.MSELoss())}

        self.trainer: Optional[Engine] = None
        self.evaluator: Optional[Engine] = None

    def set_dataset(self