saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-29 20:14:01 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Enabling softmax grad and logsoftmax grad on ORT (#12614)

Browse source 
* Enabling softmax grad and logsoftmax grad on ORT

* formatting changes

* formatting changes

* reverting changes

* Changing the OpType
This commit is contained in:
abhi-ort 2022-08-23 15:49:02 -07:00 committed by GitHub
parent cb2601c5ea
commit 73e5741a9a
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 87 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

43
orttraining/orttraining/eager/opgen/opgen/atenops.py
View file

@ -40,6 +40,46 @@ class GeluGrad(ONNXOp):
self.domain = kMSDomain
class SoftmaxGrad(ONNXOp):
"""
The operator computes the grad of softmax values:
"""
def __init__(self, dY, Y, axis=None):
super().__init__(
"SoftmaxGrad_13",
1,
[
{"at::kDouble", "at::kBFloat16", "at::kHalf", "at::kFloat"},
{"at::kDouble", "at::kBFloat16", "at::kHalf", "at::kFloat"},
],
dY,
Y,
axis=ONNXAttr(axis, AttrType.INT),
)
self.domain = kMSDomain
class LogSoftmaxGrad(ONNXOp):
"""
The operator computes the grad of log of softmax values:
"""
def __init__(self, dY, Y, axis=None):
super().__init__(
"LogSoftmaxGrad_13",
1,
[
{"at::kDouble", "at::kBFloat16", "at::kHalf", "at::kFloat"},
{"at::kDouble", "at::kBFloat16", "at::kHalf", "at::kFloat"},
],
dY,
Y,
axis=ONNXAttr(axis, AttrType.INT),
)
self.domain = kMSDomain
ops = {}
type_promotion_ops = []
aten_output_type = {}
@ -163,7 +203,8 @@ hand_implemented = {
# Leaving nll_loss_forward.output set to fallback. https://github.com/microsoft/onnxruntime/blob/main/docs/OperatorKernels.md.
"aten::nll_loss_forward.output": MakeTorchFallback(),
"aten::nll_loss_backward.grad_input": MakeTorchFallback(),
"aten::_log_softmax_backward_data.out": MakeTorchFallback(),
"aten::_softmax_backward_data": SoftmaxGrad("grad_output", "output", axis="dim"),
"aten::_log_softmax_backward_data": LogSoftmaxGrad("grad_output", "output", axis="dim"),
"aten::squeeze.dim": Squeeze("self", "dim"),
"aten::squeeze": SignatureOnly(),
"aten::unsqueeze": Unsqueeze(data="self", axes="dim"),

46
orttraining/orttraining/eager/test/ort_ops.py
View file

@ -8,7 +8,7 @@ import unittest
import numpy as np
import onnxruntime_pybind11_state as torch_ort
import torch
from parameterized import parameterized
from parameterized import parameterized, param
class OrtOpTests(unittest.TestCase):
@ -211,6 +211,50 @@ class OrtOpTests(unittest.TestCase):
assert torch.allclose(cpu_result_c, ort_result_c.cpu())
assert torch.allclose(ort_result_a.cpu(), ort_result_c.cpu())
@parameterized.expand(
[
param((2, 64, 256), 0),
param((2, 64, 256), 1),
param((2, 64, 256), -1),
param((4096, 1024), 0),
param((512, 8192), 1),
]
)
def test_softmax_grad(self, input_shape, dim):
device = self.get_device()
cpu_tensor = torch.nn.Parameter(torch.rand(input_shape))
ort_tensor = torch.nn.Parameter(cpu_tensor.detach().clone().to(device))
cpu_result = torch.softmax(cpu_tensor, dim=dim)
ort_result = torch.softmax(ort_tensor, dim=dim)
cpu_loss = cpu_result.pow(2).sum()
ort_loss = ort_result.cpu().pow(2).sum()
cpu_loss.backward()
ort_loss.backward()
assert torch.allclose(ort_result.cpu(), cpu_result)
assert torch.allclose(ort_tensor.grad.cpu(), cpu_tensor.grad, rtol=0.01)
@parameterized.expand(
[
param((2, 64, 256), 0),
param((2, 32, 128), 1),
param((2, 64, 128), -1),
param((1024, 8), 0),
param((2, 2048), 1),
]
)
def test_logsoftmax_grad(self, input_shape, dim):
device = self.get_device()
cpu_tensor = torch.nn.Parameter(torch.rand(input_shape))
ort_tensor = torch.nn.Parameter(cpu_tensor.detach().clone().to(device))
cpu_result = torch.log_softmax(cpu_tensor, dim=dim)
ort_result = torch.log_softmax(ort_tensor, dim=dim)
cpu_loss = cpu_result.pow(2).sum()
ort_loss = ort_result.cpu().pow(2).sum()
cpu_loss.backward()
ort_loss.backward()
assert torch.allclose(ort_result.cpu(), cpu_result)
assert torch.allclose(ort_tensor.grad.cpu(), cpu_tensor.grad, rtol=0.05)
def test_addmm(self):
device = self.get_device()
size = 4