class xla_trace(*args, **kwargs)[源代码]
Wraps a callable, and provides accelerated evaluation compiled by xla.

Currently it is an experimental feature. Refer to trace for more information.


import numpy as np
from basecls.models.resnet import resnet18
from megengine.autodiff.grad_manager import GradManager
from megengine.jit import xla_trace
from megengine.optimizer import Adam

model = resnet18()
gm = GradManager()
opt = Adam(model.parameters(), lr=1e-4)

# Only tensors in wrapped func args/kwargs will be treated as graph inputs,
# and other tensors will be captured as const value.
# Module, optimizer, and train data/label should be arguments of the wrapped function.
def train_step(model, opt, data, label):
    with gm:
        pred = model(data)
        loss = F.loss.cross_entropy(pred, label)
    return loss