1import pytest 2 3import torch 4 5from .fuser import set_fuser 6from .runner import get_nn_runners 7 8 9@pytest.fixture(scope="class") 10def modeldef(request, net_name, executor, fuser): 11 set_fuser(fuser, executor) 12 13 # Given a 'net_name' provided by generate_tests, build the thing 14 name, rnn_creator, context = get_nn_runners(net_name)[0] 15 creator_args = creator_args = { 16 "seqLength": 100, 17 "numLayers": 1, 18 "inputSize": 512, 19 "hiddenSize": 512, 20 "miniBatch": 64, 21 "device": "cuda", 22 "seed": None, 23 } 24 return rnn_creator(**creator_args) 25 26 27def cuda_sync(func, *args, **kwargs): 28 out = func(*args, **kwargs) 29 torch.cuda.synchronize() 30 return out 31 32 33@pytest.mark.benchmark( 34 warmup=True, 35 warmup_iterations=3, 36 disable_gc=True, 37 max_time=0.1, 38 group="fastrnns", 39) 40class TestBenchNetwork: 41 # See 'modeldef' fixture, which provides the things to benchmark 42 def test_forward(self, modeldef, benchmark): 43 forward_output = benchmark(cuda_sync, modeldef.forward, *modeldef.inputs) 44 45 def test_backward(self, modeldef, benchmark): 46 backward_input = modeldef.forward(*modeldef.inputs) 47 if modeldef.backward_setup is not None: 48 backward_input = modeldef.backward_setup(backward_input) 49 50 if modeldef.backward is not None: 51 benchmark(cuda_sync, modeldef.backward, *backward_input, retain_graph=True) 52 53 with torch.no_grad(): 54 for param in modeldef.params: 55 assert param.grad is not None 56 param.grad.zero_() 57