mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 12:21:27 +01:00
43fe45ab0f
Summary: This PR modifies `benchmarks/tensorexpr`. It follows up[ https://github.com/pytorch/pytorch/issues/44101](https://github.com/pytorch/pytorch/pull/44101) and further supports characterizing fusers with dynamic shape benchmarks. Dynamic shape condition models the use case when the input tensor shape changes in each call to the graph. Changes include: Added an auxiliary class `DynamicShape `that provides a simple API for enabling dynamic shapes in existing test cases, example can be found with `DynamicSimpleElementBench` Created new bench_cls: `DynamicSimpleElementBench`, `DynamicReduce2DInnerBench`, `DynamicReduce2DOuterBench`, and `DynamicLSTM`. They are all dynamic shaped versions of existing benchmarks and examples of enabling dynamic shape with `DynamicShape`. Pull Request resolved: https://github.com/pytorch/pytorch/pull/46107 Reviewed By: glaringlee Differential Revision: D24229400 Pulled By: bertmaher fbshipit-source-id: 889fece5ea87d0f6f6374d31dbe11b1cd1380683
106 lines
3.1 KiB
Python
106 lines
3.1 KiB
Python
from . import benchmark
|
|
import torch
|
|
|
|
class RNNEltwise(benchmark.Benchmark):
|
|
def __init__(self, mode, device, dtype, b, hs):
|
|
super().__init__(mode, device, dtype)
|
|
self.b = b
|
|
self.hs = hs
|
|
self.input = self.rand(
|
|
[b, 4 * hs], device=device, dtype=dtype, requires_grad=self.requires_grad
|
|
)
|
|
self.hx = self.rand(
|
|
[b, 4 * hs], device=device, dtype=dtype, requires_grad=self.requires_grad
|
|
)
|
|
self.cx = self.rand(
|
|
[b, hs], device=device, dtype=dtype, requires_grad=self.requires_grad
|
|
)
|
|
self.b_ih = self.rand(
|
|
[b, 4 * hs], device=device, dtype=dtype, requires_grad=self.requires_grad
|
|
)
|
|
self.b_hh = self.rand(
|
|
[b, 4 * hs], device=device, dtype=dtype, requires_grad=self.requires_grad
|
|
)
|
|
self.inputs = [
|
|
self.input,
|
|
self.hx,
|
|
self.cx,
|
|
self.b_ih,
|
|
self.b_hh,
|
|
]
|
|
|
|
def forward(self, input, hx, cx, b_ih, b_hh):
|
|
gates = input + hx + b_ih + b_hh
|
|
|
|
ingate, forgetgate, cellgate, outgate = gates.chunk(4, 1)
|
|
|
|
ingate = torch.sigmoid(ingate)
|
|
forgetgate = torch.sigmoid(forgetgate)
|
|
cellgate = torch.tanh(cellgate)
|
|
outgate = torch.sigmoid(outgate)
|
|
|
|
cy = (forgetgate * cx) + (ingate * cellgate)
|
|
hy = outgate * torch.tanh(cy)
|
|
|
|
return hy, cy
|
|
|
|
def config(self):
|
|
return [self.b, self.hs]
|
|
|
|
@staticmethod
|
|
def module():
|
|
return "rnn_eltwise"
|
|
|
|
def memory_workload(self):
|
|
def memsize(t):
|
|
return t.numel() * t.element_size()
|
|
|
|
input_size = sum([memsize(t) for t in self.inputs])
|
|
output_size = 2 * memsize(self.cx)
|
|
io_size = input_size + output_size
|
|
return {"sol": io_size, "algorithmic": io_size}
|
|
|
|
@staticmethod
|
|
def default_configs():
|
|
return [[64, 512]]
|
|
|
|
benchmark.register_benchmark_class(RNNEltwise)
|
|
|
|
|
|
class DynamicLSTM(benchmark.DynamicShape, RNNEltwise):
|
|
def __init__(self, mode, device, dtype, b, hs):
|
|
benchmark.DynamicShape.__init__(self)
|
|
RNNEltwise.__init__(self, mode, device, dtype, b, hs)
|
|
|
|
def instantiate_input(self):
|
|
b, hs = self.rand_shape([self.b, self.hs])
|
|
|
|
self.input = self.rand(
|
|
[b, 4 * hs], device=self.device, dtype=self.dtype, requires_grad=self.requires_grad
|
|
)
|
|
self.hx = self.rand(
|
|
[b, 4 * hs], device=self.device, dtype=self.dtype, requires_grad=self.requires_grad
|
|
)
|
|
self.cx = self.rand(
|
|
[b, hs], device=self.device, dtype=self.dtype, requires_grad=self.requires_grad
|
|
)
|
|
self.b_ih = self.rand(
|
|
[b, 4 * hs], device=self.device, dtype=self.dtype, requires_grad=self.requires_grad
|
|
)
|
|
self.b_hh = self.rand(
|
|
[b, 4 * hs], device=self.device, dtype=self.dtype, requires_grad=self.requires_grad
|
|
)
|
|
self.inputs = [
|
|
self.input,
|
|
self.hx,
|
|
self.cx,
|
|
self.b_ih,
|
|
self.b_hh,
|
|
]
|
|
|
|
@staticmethod
|
|
def module():
|
|
return "dynamic_lstm"
|
|
|
|
benchmark.register_benchmark_class(DynamicLSTM)
|