/
githubmirror
/
ColossalAI
Обзор
Документация
Войти
/
githubmirror
/
ColossalAI
Код
Запросы
0
Пакеты
0
Релизы
0
Аналитика
Безопасность
main
tests/kit/model_zoo/executor.py
58 строк
1 KB
Hongxin Liu
[test] merge old components to test to model zoo (#4945)
20 окт 2023, 05:35
Не верифицирован
20 окт 2023, 05:35
b8e770c
Код
Авторство
О чём код?
from typing import Callable, Dict, Optional, Union import torch from torch.nn import Module from torch.optim import Optimizer from colossalai.interface import OptimizerWrapper def run_fwd( model: Module, data: Dict, output_transform_fn: Callable, criterion: Optional[Callable] = None ) -> torch.Tensor: """run_fwd run fwd for the model Args: model (torch.nn.Module): a PyTorch model data (torch.Tensor): input data label (torch.Tensor): label criterion (Optional[Callable]): a function of criterion Returns: torch.Tensor: loss of fwd """ outputs = model(**data) outputs = output_transform_fn(outputs) if criterion: loss = criterion(outputs) else: loss = next(iter(outputs.values())).sum() return loss def run_fwd_bwd( model: Module, data: Dict, output_transform_fn: Callable, criterion: Optional[Callable] = None, optimizer: Optional[Union[Optimizer, OptimizerWrapper]] = None, ) -> torch.Tensor: """run_fwd_bwd run fwd and bwd for the model Args: model (torch.nn.Module): a PyTorch model data (torch.Tensor): input data label (torch.Tensor): label criterion (Optional[Callable]): a function of criterion Returns: torch.Tensor: loss of fwd """ loss = run_fwd(model, data, output_transform_fn, criterion) if optimizer: optimizer.backward(loss) else: loss.backward() return loss