def trace_(func, example_inputs, *args, **kwargs):
# Disable mix precision. torch.jit.trace will check the traced output
# against what is expected. Since mix precision will lead to
# loss of accuracy, this will raise warning during torch.jit.trace
orig_mixed_type = ipex.get_auto_mix_precision()
ipex.enable_auto_mix_precision(None)
jit_m = orig_trace(func, example_inputs, *args, **kwargs)
if core.get_jit_opt() and hasattr(jit_m, '_c'):
jit_m = wrap_cpp_module(torch._C._jit_pass_fold_convbn(jit_m._c))
ipex.enable_auto_mix_precision(orig_mixed_type)
return jit_m
def script_(obj, optimize=None, _frames_up=0, _rcb=None):
torch.jit.script = orig_script
jit_m = orig_script(obj,
optimize=optimize,
_frames_up=_frames_up + 1,
_rcb=_rcb)
torch.jit.script = script_
if core.get_jit_opt() and hasattr(jit_m, '_c'):
# Disable mix precision in model fusion, since mixed precision cannot
# bring any benefits for inference, but will lead to loss of accuracy
orig_mixed_type = ipex.get_auto_mix_precision()
ipex.enable_auto_mix_precision(None)
jit_m = wrap_cpp_module(torch._C._jit_pass_fold_convbn(jit_m._c))
ipex.enable_auto_mix_precision(orig_mixed_type)
return jit_m
def __init__(self, enable_or_not=False, train=False):
self.old_value = ipex.get_auto_mix_precision()
self.train_old_value = ipex.get_train()
self.enable_or_not = enable_or_not
self.train = train
def __init__(self, enable_or_not=False, train=False):
self.old_value = ipex.get_auto_mix_precision()
self.pre_running_mode = 'training' if ipex.get_train() else 'inference'
self.enable_or_not = enable_or_not
self.running_mode = 'training' if train else 'inference'