mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2025-05-03 22:18:23 +00:00
50793b35f4
* [gemini] support don't scatter after inference * [chat] update colossalai strategy * [chat] fix opt benchmark * [chat] update opt benchmark * [gemini] optimize inference * [test] add gemini inference test * [chat] fix unit test ci * [chat] fix ci * [chat] fix ci * [chat] skip checkpoint test
47 lines
1.1 KiB
Python
47 lines
1.1 KiB
Python
import torch
|
|
|
|
|
|
def run_fwd(model, data, label, criterion) -> 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
|
|
"""
|
|
if criterion:
|
|
y = model(data)
|
|
y = y.float()
|
|
loss = criterion(y, label)
|
|
else:
|
|
loss = model(data, label)
|
|
|
|
loss = loss.float()
|
|
return loss
|
|
|
|
|
|
def run_fwd_bwd(model, data, label, criterion, optimizer=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, label, criterion)
|
|
if optimizer:
|
|
optimizer.backward(loss)
|
|
else:
|
|
loss.backward()
|
|
return loss
|