def test_grad_with_multiple_inputs():
if not pytorch_pfn_extras.requires("1.8.0"):
pytest.skip('skip for PyTorch 1.7 or earlier')
if pytorch_pfn_extras.requires('1.10.0') and sys.platform == 'win32':
pytest.skip(
'ONNX grad test does not work in windows CI for torch >= 1.10')
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv = nn.Conv2d(2, 6, 3)
self.linear = nn.Linear(32, 20, bias=False)
def forward(self, x):
x0 = x * 0.5
x1 = x * 2.0
x0.requires_grad_(True)
x1.requires_grad_(True)
h = self.conv(torch.cat([x0, x1], dim=1))
grad_x = grad(
h,
(x0, x1),
retain_graph=True,
create_graph=True,
)
h = self.linear(grad_x[0])
return h
model = Net()
x = torch.ones((1, 1, 32, 32))
output_dir = _helper(
model,
x,
'grad',
enable_onnx_checker=False,
use_pfto=False,
)
actual_onnx = onnx.load(os.path.join(output_dir, 'model.onnx'))
named_nodes = {n.name: n for n in actual_onnx.graph.node}
assert 'Conv_5' in named_nodes
assert 'Gradient_6' in named_nodes
assert 'MatMul_8' in named_nodes
assert list([v.name for v in actual_onnx.graph.output]) == [
"v13_MatMul", "Gradient_y_0", "Gradient_x_0_0", "Gradient_x_1_0"
]
assert named_nodes["Concat_4"].input[0] == "Gradient_x_0_0"
assert named_nodes["Concat_4"].input[1] == "Gradient_x_1_0"
assert named_nodes["Conv_5"].output[0] == "Gradient_y_0"
def test_no_as_output():
if not pytorch_pfn_extras.requires("1.8.0"):
pytest.skip('skip for PyTorch 1.7 or earlier')
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv = nn.Conv2d(1, 6, 3)
self.linear = nn.Linear(30, 20, bias=False)
def forward(self, x):
h = self.conv(x)
h = self.linear(h)
return h
model = Net()
x = torch.ones((1, 1, 32, 32))
output_dir = _helper(model, x, 'as_output')
actual_onnx = onnx.load(os.path.join(output_dir, 'model.onnx'))
named_nodes = {n.name: n for n in actual_onnx.graph.node}
assert 'Conv_0' in named_nodes
assert 'MatMul_2' in named_nodes
assert len([v.name for v in actual_onnx.graph.output]) == 1
def optimize_onnx(self, graph: torch._C.Graph) -> torch._C.Graph:
if pytorch_pfn_extras.requires("1.9.0"):
self.run_jit_pass(torch._C._jit_pass_onnx_scalar_type_analysis, graph, self.onnx_lowprecision_cast, self.opset_version)
else:
self.run_jit_pass(torch._C._jit_pass_onnx_scalar_type_analysis, graph)
if self.do_constant_folding and self.opset_version in torch.onnx.constant_folding_opset_versions:
folded: Dict[str, torch.IValue] = torch._C._jit_pass_onnx_constant_fold( # type: ignore[attr-defined]
graph, self.vars, self.opset_version
)
# Replace input with constant nodes
input_table: Dict[str, torch._C.Value] = {i.debugName(): i for i in graph.inputs()}
for k, t in folded.items():
c: torch._C.Value = graph.create("onnx::Constant", 1).output()
assert isinstance(t, torch.Tensor)
c.node().t_("value", cast(torch.Tensor, t))
graph.prependNode(c.node())
# TODO(twata): Determine folded nodes from original graph and document it
self.node_doc_string[c.node()] = f"Constant folded node: {input_table[k]}"
input_table[k].replaceAllUsesWith(c)
c.copyMetadata(input_table[k])
self.attrs[_unique_id(c)] = ONNXValueID(k)
self.vars[k] = t
del input_table[k]
for _ in range(len(list(graph.inputs())) - len(input_table)):
graph.eraseInput(len(input_table))
torch._C._jit_pass_dce_allow_deleting_nodes_with_side_effects(graph) # type: ignore[attr-defined]
if self.onnx_peephole:
self.run_jit_pass(torch._C._jit_pass_onnx_peephole, graph, self.opset_version, self.fixed_batch_size)
return graph
def test_grad_no_export():
if not pytorch_pfn_extras.requires("1.8.0"):
pytest.skip('skip for PyTorch 1.7 or earlier')
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv = nn.Conv2d(1, 6, 3)
self.linear = nn.Linear(32, 20, bias=False)
def forward(self, x):
x = x * 0.5
x.requires_grad_(True)
h = self.conv(x)
grad_x = grad(
h,
(x, ),
retain_graph=True,
create_graph=True,
)[0]
h = self.linear(grad_x)
return h
model = Net()
x = torch.ones((1, 1, 32, 32))
y = model(x)
assert y.shape == (1, 1, 32, 20)
def test_annotate():
if not pytorch_pfn_extras.requires("1.8.0"):
pytest.skip('skip for PyTorch 1.7 or earlier')
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv = nn.Conv2d(1, 6, 3)
self.conv2 = nn.Conv2d(6, 12, 3)
self.linear = nn.Linear(28, 10, bias=False)
self.linear2 = nn.Linear(10, 5, bias=False)
def forward(self, x):
with annotate(aaa='a', bbb=['b', 'c']):
h = self.conv(x)
h = self.conv2(h)
with annotate(zzz=99, yyy=[9, 9]):
h = self.linear(h)
h = self.linear2(h)
return h
model = Net()
x = torch.ones((1, 1, 32, 32))
output_dir = _helper(model, x, 'annotate', use_pfto=False)
actual_onnx = onnx.load(os.path.join(output_dir, 'model.onnx'))
named_nodes = {n.name: n for n in actual_onnx.graph.node}
assert 'Conv_0' in named_nodes
assert 'Conv_1' in named_nodes
assert 'MatMul_3' in named_nodes
assert 'MatMul_5' in named_nodes
node_conv_0_attrs = [a.name for a in named_nodes['Conv_0'].attribute]
assert 'aaa' in node_conv_0_attrs
assert 'bbb' in node_conv_0_attrs
assert 'zzz' not in node_conv_0_attrs
assert 'yyy' not in node_conv_0_attrs
node_conv_1_attrs = [a.name for a in named_nodes['Conv_1'].attribute]
assert 'aaa' not in node_conv_1_attrs
assert 'bbb' not in node_conv_1_attrs
assert 'zzz' not in node_conv_1_attrs
assert 'yyy' not in node_conv_1_attrs
node_matmul_2_attrs = [a.name for a in named_nodes['MatMul_3'].attribute]
assert 'aaa' not in node_matmul_2_attrs
assert 'bbb' not in node_matmul_2_attrs
assert 'zzz' in node_matmul_2_attrs
assert 'yyy' in node_matmul_2_attrs
node_matmul_5_attrs = [a.name for a in named_nodes['MatMul_5'].attribute]
assert 'aaa' not in node_matmul_5_attrs
assert 'bbb' not in node_matmul_5_attrs
assert 'zzz' in node_matmul_5_attrs
assert 'yyy' in node_matmul_5_attrs
def _export_util(
model: torch.nn.Module,
args: Sequence[Any],
f: IO,
**kwargs: Any,
) -> Any:
"""Wrap operator type to export
Copied from torch.onnx.utils.export, to get output values.
"""
aten = kwargs.get('aten', False)
export_raw_ir = kwargs.get('export_raw_ir', False)
operator_export_type = kwargs.get('operator_export_type', None)
if aten or export_raw_ir:
assert operator_export_type is None
assert aten ^ export_raw_ir
# Note: OperatorExportTypes.RAW unavailable in PyTorch 1.10+
operator_export_type = OperatorExportTypes.ONNX_ATEN if\
aten else OperatorExportTypes.RAW # type: ignore
elif operator_export_type is None:
if torch.onnx.PYTORCH_ONNX_CAFFE2_BUNDLE:
operator_export_type = OperatorExportTypes.ONNX_ATEN_FALLBACK
else:
operator_export_type = OperatorExportTypes.ONNX
old_model_to_graph = torch.onnx.utils._model_to_graph
# TODO(ecastill) _model_to_graph shouldn't be direclty overriden
# This is a temporal workaround until a fix is introduced in PyTorch.
try:
torch.onnx.utils._model_to_graph = _model_to_graph_with_value_names
if pytorch_pfn_extras.requires('1.10.0'):
checker_error = getattr(torch.onnx, "CheckerError", None)
if checker_error is None:
checker_error = torch.onnx.utils.ONNXCheckerError # type: ignore[attr-defined]
try:
enable_onnx_checker = kwargs.pop('enable_onnx_checker', None)
return torch_export( # type: ignore[no-untyped-call]
model, args, f, **kwargs)
except checker_error:
if enable_onnx_checker:
raise
else:
kwargs['_retain_param_name'] = True
return torch_export( # type: ignore[no-untyped-call]
model, args, f, **kwargs)
finally:
torch.onnx.utils._model_to_graph = old_model_to_graph
def test_export_testcase_return_output():
model = nn.Sequential(nn.Linear(5, 10, bias=False))
x = torch.zeros((2, 5))
output_dir = _get_output_dir('export_filename')
if pytorch_pfn_extras.requires("1.6.0"):
with pytest.warns(UserWarning):
(out, ) = export_testcase(model, x, output_dir, return_output=True)
else:
(out, ) = export_testcase(model, x, output_dir, return_output=True)
assert os.path.isfile(os.path.join(output_dir, 'model.onnx'))
expected_out = torch.zeros((2, 10)) # check only shape size
np.testing.assert_allclose(out.detach().cpu().numpy(),
expected_out.detach().cpu().numpy())
def test_export_testcase_with_unused_input(keep_initializers_as_inputs):
if not pytorch_pfn_extras.requires("1.7.0"):
pytest.skip('skip for PyTorch 1.6 or earlier')
model = NetWithUnusedInput().to('cpu')
x = torch.zeros((1, 1, 28, 28))
unused = torch.zeros((1, ))
# Without input_names
output_dir = _helper(
model,
args=(x, unused),
d='net_with_unused_input_without_input_names',
opset_version=11,
strip_doc_string=False,
keep_initializers_as_inputs=keep_initializers_as_inputs)
assert os.path.isdir(output_dir)
test_data_set_dir = os.path.join(output_dir, 'test_data_set_0')
assert os.path.exists(os.path.join(test_data_set_dir, 'input_0.pb'))
assert not os.path.exists(os.path.join(test_data_set_dir, 'input_1.pb'))
xmodel = onnx.load_model(os.path.join(output_dir, 'model.onnx'))
assert xmodel.graph.input[0].name == 'input_0'
assert len(xmodel.graph.input) == 1 or \
xmodel.graph.input[1].name != 'input_1'
# With input_names
output_dir = _helper(
model,
args=(x, unused),
d='net_with_unused_input_with_input_names',
opset_version=11,
strip_doc_string=False,
keep_initializers_as_inputs=keep_initializers_as_inputs,
input_names=['x', 'unused'])
assert os.path.isdir(output_dir)
test_data_set_dir = os.path.join(output_dir, 'test_data_set_0')
assert os.path.exists(os.path.join(test_data_set_dir, 'input_0.pb'))
assert not os.path.exists(os.path.join(test_data_set_dir, 'input_1.pb'))
xmodel = onnx.load_model(os.path.join(output_dir, 'model.onnx'))
assert xmodel.graph.input[0].name == 'x'
assert len(xmodel.graph.input) == 1 or \
xmodel.graph.input[1].name != 'unused'
def test_as_output_no_export():
if not pytorch_pfn_extras.requires("1.8.0"):
pytest.skip('skip for PyTorch 1.7 or earlier')
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv = nn.Conv2d(1, 6, 3)
self.linear = nn.Linear(30, 20, bias=False)
def forward(self, x):
h = self.conv(x)
h = as_output("h", h)
h = self.linear(h)
return h
model = Net()
x = torch.ones((1, 1, 32, 32))
y = model(x)
assert y.shape == (1, 6, 30, 20)
def _export(
model: torch.nn.Module,
args: Sequence[Any],
strip_large_tensor_data: bool = False,
large_tensor_threshold: int = LARGE_TENSOR_DATA_THRESHOLD,
use_pfto: bool = False,
**kwargs: Any,
) -> Tuple[onnx.ModelProto, Any]:
model.zero_grad()
bytesio = io.BytesIO()
opset_ver = kwargs.get('opset_version', None)
if opset_ver is None:
opset_ver = _default_onnx_opset_version
kwargs['opset_version'] = opset_ver
if use_pfto or not pytorch_pfn_extras.requires('1.10.0'):
strip_doc_string = kwargs.get('strip_doc_string', True)
kwargs['strip_doc_string'] = False
else:
strip_doc_string = kwargs.pop('strip_doc_string', True)
kwargs['verbose'] = True
with init_annotate(model, opset_ver) as ann, \
as_output.trace(model) as (model, outputs), \
grad.init_grad_state():
if use_pfto:
outs = pfto_export(
model, args, bytesio, **kwargs)
else:
outs = _export_util(
model, args, bytesio, **kwargs)
onnx_graph = onnx.load(io.BytesIO(bytesio.getvalue()))
onnx_graph = ann.set_annotate(onnx_graph)
onnx_graph = ann.reorg_anchor(onnx_graph)
outputs.add_outputs_to_model(onnx_graph)
if strip_doc_string:
for node in onnx_graph.graph.node:
node.doc_string = b''
if strip_large_tensor_data:
_strip_large_initializer_raw_data(onnx_graph, large_tensor_threshold)
return onnx_graph, outs
class TestDistributedDataParallel:
def test_save_load(self):
module = MyModule()
with_ddp = DistributedDataParallel(module)
assert module.state_dict().keys() == with_ddp.state_dict().keys()
module.load_state_dict(with_ddp.state_dict())
assert np.array_equal(module.state_dict()["param0"],
with_ddp.state_dict()["param0"])
assert np.array_equal(module.state_dict()["param1"],
with_ddp.state_dict()["param1"])
assert np.array_equal(module.state_dict()["buffer"],
with_ddp.state_dict()["buffer"])
@pytest.mark.parametrize('device_type', _device_types())
def test_sync_init_params(self, device_type):
module0 = MyModule()
module0.param0.data = torch.tensor([1.])
r0, r1 = _launch(inputs=[torch.tensor([1.]),
torch.tensor([2.])],
modules=[module0, MyModule()],
device_type=device_type)
assert r0[0].item() == 1
assert r1[0].item() == 2
assert r0[1]["param0"].item() == 1.0
assert r1[1]["param0"].item() == 1.0
@pytest.mark.parametrize('device_type', _device_types())
def test_all_reduce(self, device_type):
r0, r1 = _launch(inputs=[torch.tensor([1.]),
torch.tensor([2.])],
device_type=device_type)
assert r0[0].item() == -1
assert r1[0].item() == -2
assert r0[2]["module.param0"].item() == 1.5
assert r1[2]["module.param0"].item() == 1.5
assert r0[2]["module.param1"] is None
assert r1[2]["module.param1"] is None
@pytest.mark.parametrize('device_type', _device_types())
def test_specific_reduce(self, device_type):
r0, r1 = _launch(inputs=[torch.tensor([1.]),
torch.tensor([2.])],
args={"reduce_function": Collectives._to_zero},
device_type=device_type)
assert r0[2]["module.param0"].item() == 0.0
assert r1[2]["module.param0"].item() == 0.0
@pytest.mark.parametrize('device_type', _device_types())
def test_nosync_buffer(self, device_type):
r0, r1 = _launch(inputs=[torch.tensor([1.]),
torch.tensor([2.])],
args={"broadcast_buffers": False},
device_type=device_type)
assert r0[0].item() == -1
assert r1[0].item() == -2
assert r0[1]["buffer"].item() == 1
assert r1[1]["buffer"].item() == 2
@pytest.mark.parametrize('device_type', _device_types())
def test_sync_buffer(self, device_type):
r0, r1 = _launch(inputs=[torch.tensor([1.]),
torch.tensor([2.])],
args={"broadcast_buffers": True},
device_type=device_type)
assert r0[0].item() == -1
assert r1[0].item() == -2
assert r0[1]["buffer"].item() == 1
assert r1[1]["buffer"].item() == 1
@pytest.mark.parametrize('device_type', _device_types())
def test_specific_broadcast(self, device_type):
r0, r1 = _launch(inputs=[torch.tensor([1.]),
torch.tensor([2.])],
args={
"broadcast_function": Collectives._to_zero,
"broadcast_buffers": True
},
device_type=device_type)
assert r0[1]["buffer"].item() == 0.0
assert r1[1]["buffer"].item() == 0.0
@pytest.mark.parametrize('device_type', _device_types())
def test_define_by_run(self, device_type):
r0, r1 = _launch(inputs=[torch.tensor([1.]),
torch.tensor([-1])],
device_type=device_type)
assert r0[0].item() == -1
assert r1[0].item() == 1
assert r0[2]["module.param0"].item() == 0.5
assert r1[2]["module.param0"].item() == 0.5
assert r0[2]["module.param1"].item() == 0.5
assert r1[2]["module.param1"].item() == 0.5
@pytest.mark.parametrize('device_type', _device_types())
def test_no_sync(self, device_type):
r0, r1 = _launch(inputs=[torch.tensor([1.]),
torch.tensor([2.])],
step=Steps._step_with_no_sync,
device_type=device_type)
assert r0[0].item() == -1
assert r1[0].item() == -2
assert r0[2]["module.param0"].item() == 1
assert r1[2]["module.param0"].item() == 2
assert r0[2]["module.param1"] is None
assert r1[2]["module.param1"] is None
@pytest.mark.parametrize('device_type', _device_types())
def test_hook(self, device_type):
r0, r1 = _launch(inputs=[torch.tensor([1.]),
torch.tensor([2.])],
step=Steps._step_with_hook,
device_type=device_type)
assert r0[0].item() == -1
assert r1[0].item() == -2
assert r0[2]["module.param0"].item() == 0
assert r1[2]["module.param0"].item() == 0
assert r0[2]["module.param1"].item() == 0
assert r1[2]["module.param1"].item() == 0
@pytest.mark.parametrize('device_type', _device_types())
@pytest.mark.skipif(
not pytorch_pfn_extras.requires("1.6.0"),
reason="Variable._execution_engine.queue_callback does not work "
"with checkpointing when torch < 1.6.0")
def test_checkpoint(self, device_type):
r0, r1 = _launch(
inputs=[torch.tensor([[1.]]),
torch.tensor([[2.]])],
modules=[MyModuleWithCheckpoint(),
MyModuleWithCheckpoint()],
step=Steps._step_with_hook,
device_type=device_type)
grad0 = r0[2]
grad1 = r1[2]
for key in grad0.keys():
assert np.array_equal(grad0[key].cpu().numpy(),
grad1[key].cpu().numpy())
import os
import pytest
import torch
import pytorch_pfn_extras as ppe
_profiler_available = (os.name != 'nt' or ppe.requires("1.9"))
@pytest.mark.skipif(not _profiler_available,
reason="profiler is not available")
@pytest.mark.parametrize('device', ['cpu', 'cuda'])
def test_record(device):
if not torch.cuda.is_available() and device == 'cuda':
pytest.skip()
model = torch.nn.Linear(30, 40)
model.to(device)
x = torch.arange(30, dtype=torch.float32).to(device)
with torch.profiler.profile() as prof:
with ppe.profiler.record('my_tag_1'):
model(x)
keys = [event.key for event in prof.key_averages()]
assert 'my_tag_1' in keys
assert 'aten::linear' in keys
@pytest.mark.skipif(not _profiler_available,
reason="profiler is not available")
def test_export_testcase_strip_large_tensor_data():
if not pytorch_pfn_extras.requires("1.6.0"):
pytest.skip('skip for PyTorch 1.5 or earlier')
model = Net().to('cpu')
x = torch.zeros((1, 1, 28, 28))
output_dir = _helper(model,
x,
'mnist_stripped_tensor_data',
output_grad=True,
strip_large_tensor_data=True,
metadata=True)
assert os.path.isdir(output_dir)
assert os.path.isfile(os.path.join(output_dir, 'meta.json'))
assert os.path.isfile(os.path.join(output_dir, 'model.onnx'))
test_data_set_dir = os.path.join(output_dir, 'test_data_set_0')
assert os.path.isfile(os.path.join(test_data_set_dir, 'input_0.pb'))
assert os.path.isfile(os.path.join(test_data_set_dir, 'output_0.pb'))
for i in range(8):
assert os.path.isfile(
os.path.join(test_data_set_dir, 'gradient_{}.pb'.format(i)))
assert not os.path.isfile(os.path.join(test_data_set_dir, 'gradient_8.pb'))
with open(os.path.join(output_dir, 'meta.json')) as metaf:
metaj = json.load(metaf)
assert metaj['strip_large_tensor_data']
def is_stripped_with_check(tensor):
if is_large_tensor(tensor, LARGE_TENSOR_DATA_THRESHOLD):
assert tensor.data_location == onnx.TensorProto.EXTERNAL
assert tensor.external_data[0].key == 'location'
meta = json.loads(tensor.external_data[0].value)
assert meta['type'] == 'stripped'
assert type(meta['average']) == float
assert type(meta['variance']) == float
return True
assert len(tensor.external_data) == 0
return False
onnx_model = onnx.load(os.path.join(output_dir, 'model.onnx'),
load_external_data=False)
check_stripped = [
is_stripped_with_check(init) for init in onnx_model.graph.initializer
]
# this testcase tests strip, so output mode is no stripped, test is failed
assert any(check_stripped)
for pb_filepath in ('input_0.pb', 'output_0.pb'):
with open(os.path.join(test_data_set_dir, pb_filepath), 'rb') as f:
tensor = onnx.TensorProto()
tensor.ParseFromString(f.read())
is_stripped_with_check(tensor)
# check re-load stripped onnx
_strip_large_tensor_tool_impl(os.path.join(output_dir, 'model.onnx'),
os.path.join(output_dir, 'model_re.onnx'),
LARGE_TENSOR_DATA_THRESHOLD)
assert os.path.isfile(os.path.join(output_dir, 'model_re.onnx'))
# loading check
onnx.load(os.path.join(output_dir, 'model_re.onnx'),
load_external_data=False)
# unstrip test
unstrip_output_dir = _get_output_dir('mnist_unstripped_tensor_data')
unstrip(output_dir, out_path=unstrip_output_dir)
def is_unstripped_with_check(tensor):
if is_large_tensor(tensor, LARGE_TENSOR_DATA_THRESHOLD):
assert tensor.data_location == onnx.TensorProto.DEFAULT
return True
return False
onnx_paths = Path(unstrip_output_dir).glob('*.onnx')
check_unstripped = []
for onnx_path in onnx_paths:
onnx_model = onnx.load(onnx_path, load_external_data=False)
check_unstripped.extend([
is_unstripped_with_check(init)
for init in onnx_model.graph.initializer
])
assert len(check_unstripped) > 0
assert any(check_unstripped)
pb_paths = Path(unstrip_output_dir).glob('**/*.pb')
checked = False
for pb_path in pb_paths:
tensor = onnx.TensorProto()
with open(pb_path, 'rb') as f:
tensor.ParseFromString(f.read())
checked = is_unstripped_with_check(tensor) or checked
assert checked, 'more than one data is unstripped'
with open(os.path.join(unstrip_output_dir, 'meta.json')) as metaf:
metaj = json.load(metaf)
assert not metaj['strip_large_tensor_data']
trainer = engine.create_trainer(model_with_loss, {
"0": optimizer0,
"1": optimizer1
},
20,
device=device,
evaluator=evaluator,
extensions=extensions,
out_dir=path,
logic=ppe.handler.CodeBlockLogic())
trainer.run(data, data)
@pytest.mark.skipif(
os.name == 'nt' and not ppe.requires("1.9"),
reason='torch.profiler.profile is not supported.',
)
def test_trainer_profile():
device = 'cpu'
model = MyModel()
model_with_loss = MyModelWithLossDictOutput(model)
ppe.to(model_with_loss, device)
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
data = torch.utils.data.DataLoader([{
'x': torch.rand(20, ),
't': torch.rand(10, )
} for i in range(10)])
extensions = _make_extensions()
evaluator = engine.create_evaluator(model_with_loss, device=device)
