mirror of https://github.com/hpcaitech/ColossalAI
31 lines
800 B
Python
31 lines
800 B
Python
import torch
|
|
|
|
|
|
def run_fwd_bwd(model, data, label, criterion, use_init_ctx=False) -> 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
|
|
use_init_ctx (bool, optional): whether the model is initialized under the contxt of ColoInitCtx. Defaults to False.
|
|
|
|
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()
|
|
if use_init_ctx:
|
|
model.backward(loss)
|
|
else:
|
|
loss.backward()
|
|
return loss
|