mirror of
https://github.com/wassname/pytorch-lightning.git
synced 2026-08-20 12:40:52 +08:00
Nic Eggert
019f612204
* Run AMP tests in their own process With opt_level="O1" (the default), AMP patches many torch functions, which breaks any tests that run afterwards. This patch introduces a pytest extension that lets tests be marked with @pytest.mark.spawn so that they are run in their own process using torch.multiprocessing.spawn so that the main python interpreter stays un-patched. Note that tests using DDP already run AMP in its own process, so they don't need this annotation. * Fix AMP tests Since AMP defaults to O1 now, DP tests no longer throw exceptions. Since AMP patches torch functions, CPU inference no longer works. Skip prediction step for AMP tests. * typo
23 lines
584 B
Python
23 lines
584 B
Python
import pytest
|
|
|
|
import torch.multiprocessing as mp
|
|
|
|
|
|
def pytest_configure(config):
|
|
config.addinivalue_line("markers", "spawn: spawn test in a separate process using torch.multiprocessing.spawn")
|
|
|
|
|
|
def wrap(i, fn, args):
|
|
return fn(*args)
|
|
|
|
|
|
@pytest.mark.tryfirst
|
|
def pytest_pyfunc_call(pyfuncitem):
|
|
if pyfuncitem.get_closest_marker("spawn"):
|
|
testfunction = pyfuncitem.obj
|
|
funcargs = pyfuncitem.funcargs
|
|
testargs = tuple([funcargs[arg] for arg in pyfuncitem._fixtureinfo.argnames])
|
|
|
|
mp.spawn(wrap, (testfunction, testargs))
|
|
return True
|