mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
dd3a77bc96
Signed-off-by: Edward Z. Yang <ezyang@meta.com> Pull Request resolved: https://github.com/pytorch/pytorch/pull/105928 Approved by: https://github.com/albanD
109 lines
3.0 KiB
Python
109 lines
3.0 KiB
Python
from collections import namedtuple
|
|
from functools import partial
|
|
|
|
import torch
|
|
import torchvision.models as cnn
|
|
|
|
from .factory import (
|
|
dropoutlstm_creator,
|
|
imagenet_cnn_creator,
|
|
layernorm_pytorch_lstm_creator,
|
|
lnlstm_creator,
|
|
lstm_creator,
|
|
lstm_multilayer_creator,
|
|
lstm_premul_bias_creator,
|
|
lstm_premul_creator,
|
|
lstm_simple_creator,
|
|
pytorch_lstm_creator,
|
|
varlen_lstm_creator,
|
|
varlen_pytorch_lstm_creator,
|
|
)
|
|
|
|
|
|
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
|
|
),
|
|
}
|