mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-06 12:20:52 +01:00
bdf7cb9d9c
This reverts commit e20c9bf288.
Reverted https://github.com/pytorch/pytorch/pull/165410 on behalf of https://github.com/clee2000 due to sorry I'm going to revert this since I want to try to back out some other things that are conflicting with this, there is nothing wrong with this PR, rebasing and resolving the merge conflicts should be enough, sorry for the churn ([comment](https://github.com/pytorch/pytorch/pull/165410#issuecomment-3427532373))
235 lines
7.7 KiB
Python
235 lines
7.7 KiB
Python
# mypy: allow-untyped-defs
|
|
import torch
|
|
|
|
|
|
class MkldnnLinear(torch.jit.ScriptModule):
|
|
def __init__(self, dense_module, dtype):
|
|
super().__init__()
|
|
self.register_buffer('weight', dense_module.weight.to_mkldnn(dtype))
|
|
if dense_module.bias is not None:
|
|
# Bias can be fp32 or bf16 for OneDNN bf16 path, but for good accuracy,
|
|
# we use fp32 dtype.
|
|
self.register_buffer('bias', dense_module.bias.to_mkldnn())
|
|
else:
|
|
# TODO: Remove this once ScriptModule supports registering None buffer
|
|
self.register_buffer(
|
|
'bias',
|
|
torch.zeros([dense_module.weight.size(0)], dtype=torch.float).to_mkldnn())
|
|
|
|
@torch.jit.script_method
|
|
def __getstate__(self):
|
|
return (self.weight.to_dense(), self.bias.to_dense(), self.training)
|
|
|
|
@torch.jit.script_method
|
|
def __setstate__(self, state):
|
|
self.weight = state[0].to_mkldnn()
|
|
self.bias = state[1].to_mkldnn()
|
|
self.training = state[2]
|
|
|
|
@torch.jit.script_method
|
|
def forward(self, x):
|
|
x_mkldnn = x if x.is_mkldnn else x.to_mkldnn()
|
|
y_mkldnn = torch._C._nn.mkldnn_linear(x_mkldnn, self.weight, self.bias)
|
|
y = y_mkldnn if x.is_mkldnn else y_mkldnn.to_dense()
|
|
return y
|
|
|
|
|
|
class _MkldnnConvNd(torch.jit.ScriptModule):
|
|
"""Common base of MkldnnConv1d and MkldnnConv2d."""
|
|
|
|
__constants__ = ['stride', 'padding', 'dilation', 'groups']
|
|
|
|
def __init__(self, dense_module):
|
|
super().__init__()
|
|
|
|
self.stride = dense_module.stride
|
|
self.padding = dense_module.padding
|
|
self.dilation = dense_module.dilation
|
|
self.groups = dense_module.groups
|
|
|
|
if dense_module.bias is not None:
|
|
self.register_buffer('bias', dense_module.bias.to_mkldnn())
|
|
else:
|
|
# Bias can be fp32 or bf16 for OneDNN bf16 path, but for good accuracy,
|
|
# we use fp32 dtype.
|
|
# TODO: Remove this once ScriptModule supports registering None buffer
|
|
self.register_buffer(
|
|
'bias',
|
|
torch.zeros([dense_module.weight.size(0)], dtype=torch.float).to_mkldnn())
|
|
|
|
@torch.jit.script_method
|
|
def __getstate__(self):
|
|
return (self.weight.to_dense(), self.bias.to_dense(), self.training)
|
|
|
|
@torch.jit.script_method
|
|
def forward(self, x):
|
|
return torch.mkldnn_convolution(
|
|
x,
|
|
self.weight,
|
|
self.bias,
|
|
self.padding,
|
|
self.stride,
|
|
self.dilation,
|
|
self.groups)
|
|
|
|
|
|
class MkldnnConv1d(_MkldnnConvNd):
|
|
def __init__(self, dense_module, dtype):
|
|
super().__init__(dense_module)
|
|
|
|
self.register_buffer('weight', dense_module.weight.to_mkldnn(dtype))
|
|
|
|
@torch.jit.script_method
|
|
def __setstate__(self, state):
|
|
self.weight = state[0].to_mkldnn()
|
|
self.bias = state[1].to_mkldnn()
|
|
self.training = state[2]
|
|
|
|
|
|
class MkldnnConv2d(_MkldnnConvNd):
|
|
def __init__(self, dense_module, dtype):
|
|
super().__init__(dense_module)
|
|
|
|
self.register_buffer('weight', torch._C._nn.mkldnn_reorder_conv2d_weight(
|
|
dense_module.weight.to_mkldnn(dtype),
|
|
self.padding,
|
|
self.stride,
|
|
self.dilation,
|
|
self.groups))
|
|
|
|
@torch.jit.script_method
|
|
def __setstate__(self, state):
|
|
self.weight = torch._C._nn.mkldnn_reorder_conv2d_weight(
|
|
state[0].to_mkldnn(),
|
|
self.padding,
|
|
self.stride,
|
|
self.dilation,
|
|
self.groups)
|
|
self.bias = state[1].to_mkldnn()
|
|
self.training = state[2]
|
|
|
|
class MkldnnConv3d(_MkldnnConvNd):
|
|
def __init__(self, dense_module, dtype):
|
|
super().__init__(dense_module)
|
|
|
|
self.register_buffer('weight', torch._C._nn.mkldnn_reorder_conv3d_weight(
|
|
dense_module.weight.to_mkldnn(dtype),
|
|
self.padding,
|
|
self.stride,
|
|
self.dilation,
|
|
self.groups))
|
|
|
|
@torch.jit.script_method
|
|
def __setstate__(self, state):
|
|
self.weight = torch._C._nn.mkldnn_reorder_conv3d_weight(
|
|
state[0].to_mkldnn(),
|
|
self.padding,
|
|
self.stride,
|
|
self.dilation,
|
|
self.groups)
|
|
self.bias = state[1].to_mkldnn()
|
|
self.training = state[2]
|
|
|
|
|
|
class MkldnnBatchNorm(torch.jit.ScriptModule):
|
|
__constants__ = ['exponential_average_factor', 'eps']
|
|
|
|
def __init__(self, dense_module):
|
|
super().__init__()
|
|
|
|
assert not dense_module.training
|
|
assert dense_module.track_running_stats
|
|
assert dense_module.affine
|
|
|
|
if dense_module.momentum is None:
|
|
self.exponential_average_factor = 0.0
|
|
else:
|
|
self.exponential_average_factor = dense_module.momentum
|
|
self.eps = dense_module.eps
|
|
|
|
self.register_buffer('weight', dense_module.weight.to_mkldnn())
|
|
self.register_buffer('bias', dense_module.bias.to_mkldnn())
|
|
self.register_buffer('running_mean', dense_module.running_mean.to_mkldnn())
|
|
self.register_buffer('running_var', dense_module.running_var.to_mkldnn())
|
|
|
|
@torch.jit.script_method
|
|
def __getstate__(self):
|
|
weight = self.weight.to_dense()
|
|
bias = self.bias.to_dense()
|
|
running_mean = self.running_mean.to_dense()
|
|
running_var = self.running_var.to_dense()
|
|
return (weight, bias, running_mean, running_var, self.training)
|
|
|
|
@torch.jit.script_method
|
|
def __setstate__(self, state):
|
|
self.weight = state[0].to_mkldnn()
|
|
self.bias = state[1].to_mkldnn()
|
|
self.running_mean = state[2].to_mkldnn()
|
|
self.running_var = state[3].to_mkldnn()
|
|
self.training = state[4]
|
|
|
|
@torch.jit.script_method
|
|
def forward(self, x):
|
|
return torch.batch_norm(
|
|
x,
|
|
self.weight,
|
|
self.bias,
|
|
self.running_mean,
|
|
self.running_var,
|
|
False, # training
|
|
self.exponential_average_factor,
|
|
self.eps,
|
|
False, # cuda_enabled
|
|
)
|
|
|
|
class MkldnnPrelu(torch.jit.ScriptModule):
|
|
def __init__(self, dense_module, dtype):
|
|
super().__init__()
|
|
self.register_buffer('weight', dense_module.weight.to_mkldnn(dtype))
|
|
|
|
@torch.jit.script_method
|
|
def __getstate__(self):
|
|
return (self.weight.to_dense(), self.training)
|
|
|
|
@torch.jit.script_method
|
|
def __setstate__(self, state):
|
|
self.weight = state[0].to_mkldnn()
|
|
self.training = state[1]
|
|
|
|
@torch.jit.script_method
|
|
def forward(self, x):
|
|
x_mkldnn = x if x.is_mkldnn else x.to_mkldnn()
|
|
y_mkldnn = torch.prelu(x_mkldnn, self.weight)
|
|
y = y_mkldnn if x.is_mkldnn else y_mkldnn.to_dense()
|
|
return y
|
|
|
|
def to_mkldnn(module, dtype=torch.float):
|
|
assert dtype in [torch.float, torch.bfloat16, torch.half], \
|
|
"MKLDNN only support float, bfloat16, and half path now"
|
|
|
|
def m_fn(m, d):
|
|
if isinstance(m, torch.nn.Linear):
|
|
return MkldnnLinear(m, d)
|
|
elif isinstance(m, torch.nn.Conv1d):
|
|
return MkldnnConv1d(m, d)
|
|
elif isinstance(m, torch.nn.Conv2d):
|
|
return MkldnnConv2d(m, d)
|
|
elif isinstance(m, torch.nn.Conv3d):
|
|
return MkldnnConv3d(m, d)
|
|
elif isinstance(m, (torch.nn.BatchNorm2d, torch.nn.BatchNorm3d)):
|
|
# For batchnorm bf16 path, OneDNN requires weight and bias need fp32 dtype.
|
|
# so it doesn't need dtype argument.
|
|
return MkldnnBatchNorm(m)
|
|
elif isinstance(m, torch.nn.PReLU):
|
|
return MkldnnPrelu(m, d)
|
|
else:
|
|
return m
|
|
|
|
def m_fn_rec(m, d):
|
|
new_m = m_fn(m, d)
|
|
for name, sub_m in m.named_children():
|
|
setattr(new_m, name, m_fn_rec(sub_m, d))
|
|
return new_m
|
|
|
|
return m_fn_rec(module, dtype)
|