mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
3875e1ba45
Summary: Sometimes at::cat gets transposed inputs and goes on a slow path. Also, make jit_premul lstm benchmark add bias to the whole input tensor to avoid separate reduction kernels in the backward pass. Pull Request resolved: https://github.com/pytorch/pytorch/pull/18816 Differential Revision: D15013576 Pulled By: wanchaol fbshipit-source-id: bcfa1cf44180b11b05b0f55f034707012f66281a
69 lines
2.6 KiB
Python
69 lines
2.6 KiB
Python
from collections import namedtuple
|
|
from functools import partial
|
|
import torch
|
|
import torchvision.models as cnn
|
|
|
|
from .factory import *
|
|
|
|
|
|
class DisableCuDNN():
|
|
def __enter__(self):
|
|
self.saved = torch.backends.cudnn.enabled
|
|
torch.backends.cudnn.enabled = False
|
|
|
|
def __exit__(self, *args, **kwargs):
|
|
torch.backends.cudnn.enabled = self.saved
|
|
|
|
|
|
class DummyContext():
|
|
def __enter__(self):
|
|
pass
|
|
|
|
def __exit__(self, *args, **kwargs):
|
|
pass
|
|
|
|
|
|
class AssertNoJIT():
|
|
def __enter__(self):
|
|
import os
|
|
enabled = os.environ.get('PYTORCH_JIT', 1)
|
|
assert not enabled
|
|
|
|
def __exit__(self, *args, **kwargs):
|
|
pass
|
|
|
|
|
|
RNNRunner = namedtuple('RNNRunner', [
|
|
'name', 'creator', 'context',
|
|
])
|
|
|
|
|
|
def get_nn_runners(*names):
|
|
return [nn_runners[name] for name in names]
|
|
|
|
|
|
nn_runners = {
|
|
'cudnn': RNNRunner('cudnn', pytorch_lstm_creator, DummyContext),
|
|
'cudnn_dropout': RNNRunner('cudnn_dropout', partial(pytorch_lstm_creator, dropout=0.4), DummyContext),
|
|
'cudnn_layernorm': RNNRunner('cudnn_layernorm', layernorm_pytorch_lstm_creator, DummyContext),
|
|
'vl_cudnn': RNNRunner('vl_cudnn', varlen_pytorch_lstm_creator, DummyContext),
|
|
'vl_jit': RNNRunner('vl_jit', partial(varlen_lstm_creator, script=True), DummyContext),
|
|
'vl_py': RNNRunner('vl_py', varlen_lstm_creator, DummyContext),
|
|
'aten': RNNRunner('aten', pytorch_lstm_creator, DisableCuDNN),
|
|
'jit': RNNRunner('jit', lstm_creator, DummyContext),
|
|
'jit_premul': RNNRunner('jit_premul', lstm_premul_creator, DummyContext),
|
|
'jit_premul_bias': RNNRunner('jit_premul_bias', lstm_premul_bias_creator, DummyContext),
|
|
'jit_simple': RNNRunner('jit_simple', lstm_simple_creator, DummyContext),
|
|
'jit_multilayer': RNNRunner('jit_multilayer', lstm_multilayer_creator, DummyContext),
|
|
'jit_layernorm': RNNRunner('jit_layernorm', lnlstm_creator, DummyContext),
|
|
'jit_layernorm_decom': RNNRunner('jit_layernorm_decom',
|
|
partial(lnlstm_creator, decompose_layernorm=True),
|
|
DummyContext),
|
|
'jit_dropout': RNNRunner('jit_dropout', dropoutlstm_creator, DummyContext),
|
|
'py': RNNRunner('py', partial(lstm_creator, script=False), DummyContext),
|
|
'resnet18': RNNRunner('resnet18', imagenet_cnn_creator(cnn.resnet18, jit=False), DummyContext),
|
|
'resnet18_jit': RNNRunner('resnet18_jit', imagenet_cnn_creator(cnn.resnet18), DummyContext),
|
|
'resnet50': RNNRunner('resnet50', imagenet_cnn_creator(cnn.resnet50, jit=False), DummyContext),
|
|
'resnet50_jit': RNNRunner('resnet50_jit', imagenet_cnn_creator(cnn.resnet50), DummyContext),
|
|
}
|