import numpy
import pytest
from thinc.api import (
SGD,
ArgsKwargs,
CupyOps,
Linear,
MPSOps,
NumpyOps,
PyTorchWrapper,
PyTorchWrapper_v2,
PyTorchWrapper_v3,
Relu,
chain,
get_current_ops,
torch2xp,
use_ops,
xp2torch,
)
from thinc.backends import context_pools
from thinc.compat import has_cupy_gpu, has_torch, has_torch_amp, has_torch_mps_gpu
from thinc.layers.pytorchwrapper import PyTorchWrapper_v3
from thinc.shims.pytorch import (
default_deserialize_torch_model,
default_serialize_torch_model,
)
from thinc.shims.pytorch_grad_scaler import PyTorchGradScaler
from thinc.util import get_torch_default_device
from ..util import check_input_converters, make_tempdir
XP_OPS = [NumpyOps()]
if has_cupy_gpu:
XP_OPS.append(CupyOps())
if has_torch_mps_gpu:
XP_OPS.append(MPSOps())
if has_torch_amp:
TORCH_MIXED_PRECISION = [False, True]
else:
TORCH_MIXED_PRECISION = [False]
XP_OPS_MIXED = [
(ops, mixed)
for ops in XP_OPS
for mixed in TORCH_MIXED_PRECISION
if not mixed or isinstance(ops, CupyOps)
]
def check_learns_zero_output(model, sgd, X, Y):
"""Check we can learn to output a zero vector"""
Yh, get_dX = model.begin_update(X)
dYh = (Yh - Y) / Yh.shape[0]
dX = get_dX(dYh)
model.finish_update(sgd)
prev = numpy.abs(Yh.sum())
for i in range(100):
Yh, get_dX = model.begin_update(X)
total = numpy.abs(Yh.sum())
dX = get_dX(Yh - Y) # noqa: F841
model.finish_update(sgd)
assert total < prev
@pytest.mark.skipif(not has_torch, reason="needs PyTorch")
@pytest.mark.parametrize("nN,nI,nO", [(2, 3, 4)])
def test_pytorch_unwrapped(nN, nI, nO):
model = Linear(nO, nI).initialize()
X = numpy.zeros((nN, nI), dtype="f")
X += numpy.random.uniform(size=X.size).reshape(X.shape)
sgd = SGD(0.01)
Y = numpy.zeros((nN, nO), dtype="f")
check_learns_zero_output(model, sgd, X, Y)
@pytest.mark.skipif(not has_torch, reason="needs PyTorch")
@pytest.mark.parametrize("nN,nI,nO", [(2, 3, 4)])
def test_pytorch_wrapper(nN, nI, nO):
import torch.nn
model = PyTorchWrapper(torch.nn.Linear(nI, nO)).initialize()
sgd = SGD(0.001)
X = numpy.zeros((nN, nI), dtype="f")
X += numpy.random.uniform(size=X.size).reshape(X.shape)
Y = numpy.zeros((nN, nO), dtype="f")
Yh, get_dX = model.begin_update(X)
assert isinstance(Yh, numpy.ndarray)
assert Yh.shape == (nN, nO)
dYh = (Yh - Y) / Yh.shape[0]
dX = get_dX(dYh)
model.finish_update(sgd)
assert dX.shape == (nN, nI)
check_learns_zero_output(model, sgd, X, Y)
assert isinstance(model.predict(X), numpy.ndarray)
@pytest.mark.skipif(not has_torch, reason="needs PyTorch")
@pytest.mark.parametrize("ops_mixed", XP_OPS_MIXED)
@pytest.mark.parametrize("nN,nI,nO", [(2, 3, 4)])
def test_pytorch_wrapper_thinc_input(ops_mixed, nN, nI, nO):
import torch.nn
ops, mixed_precision = ops_mixed
with use_ops(ops.name):
ops = get_current_ops()
pytorch_layer = torch.nn.Linear(nO, nO)
# Initialize with large weights to trigger overflow of FP16 in
# mixed-precision training.
torch.nn.init.uniform_(pytorch_layer.weight, 9.0, 11.0)
device = get_torch_default_device()
model = chain(
Relu(),
PyTorchWrapper_v2(
pytorch_layer.to(device),
mixed_precision=mixed_precision,
grad_scaler=PyTorchGradScaler(
enabled=mixed_precision, init_scale=2.0**16
),
).initialize(),
)
# pytorch allocator is set in PyTorchShim
if isinstance(ops, CupyOps):
assert "pytorch" in context_pools.get()
sgd = SGD(0.001)
X = ops.xp.zeros((nN, nI), dtype="f")
X += ops.xp.random.uniform(size=X.size).reshape(X.shape)
Y = ops.xp.zeros((nN, nO), dtype="f")
model.initialize(X, Y)
Yh, get_dX = model.begin_update(X)
assert isinstance(Yh, ops.xp.ndarray)
assert Yh.shape == (nN, nO)
dYh = (Yh - Y) / Yh.shape[0]
dX = get_dX(dYh)
model.finish_update(sgd)
assert dX.shape == (nN, nI)
check_learns_zero_output(model, sgd, X, Y)
assert isinstance(model.predict(X), ops.xp.ndarray)
@pytest.mark.skipif(not has_torch, reason="needs PyTorch")
def test_pytorch_roundtrip_conversion():
import torch
xp_tensor = numpy.zeros((2, 3), dtype="f")
torch_tensor = xp2torch(xp_tensor)
assert isinstance(torch_tensor, torch.Tensor)
new_xp_tensor = torch2xp(torch_tensor)
assert numpy.array_equal(xp_tensor, new_xp_tensor)
@pytest.mark.skipif(not has_torch, reason="needs PyTorch")
def test_pytorch_wrapper_roundtrip():
import torch.nn
model = PyTorchWrapper(torch.nn.Linear(2, 3))
model_bytes = model.to_bytes()
PyTorchWrapper(torch.nn.Linear(2, 3)).from_bytes(model_bytes)
with make_tempdir() as path:
model_path = path / "model"
model.to_disk(model_path)
new_model = PyTorchWrapper(torch.nn.Linear(2, 3)).from_bytes(model_bytes)
new_model.from_disk(model_path)
@pytest.mark.skipif(not has_torch, reason="needs PyTorch")
@pytest.mark.parametrize(
"data,n_args,kwargs_keys",
[
# fmt: off
(numpy.zeros((2, 3), dtype="f"), 1, []),
([numpy.zeros((2, 3), dtype="f"), numpy.zeros((2, 3), dtype="f")], 2, []),
((numpy.zeros((2, 3), dtype="f"), numpy.zeros((2, 3), dtype="f")), 2, []),
({"a": numpy.zeros((2, 3), dtype="f"), "b": numpy.zeros((2, 3), dtype="f")}, 0, ["a", "b"]),
(ArgsKwargs((numpy.zeros((2, 3), dtype="f"), numpy.zeros((2, 3), dtype="f")), {"c": numpy.zeros((2, 3), dtype="f")}), 2, ["c"]),
# fmt: on
],
)
def test_pytorch_convert_inputs(data, n_args, kwargs_keys):
import torch.nn
model = PyTorchWrapper(torch.nn.Linear(3, 4))
convert_inputs = model.attrs["convert_inputs"]
Y, backprop = convert_inputs(model, data, is_train=True)
check_input_converters(Y, backprop, data, n_args, kwargs_keys, torch.Tensor)
@pytest.mark.skipif(not has_torch, reason="needs PyTorch")
def test_pytorch_wrapper_custom_serde():
import torch.nn
def serialize(model):
return default_serialize_torch_model(model)
def deserialize(model, state_bytes, device):
return default_deserialize_torch_model(model, state_bytes, device)
def get_model():
return PyTorchWrapper_v3(
torch.nn.Linear(2, 3),
serialize_model=serialize,
deserialize_model=deserialize,
)
model = get_model()
model_bytes = model.to_bytes()
get_model().from_bytes(model_bytes)
with make_tempdir() as path:
model_path = path / "model"
model.to_disk(model_path)
new_model = get_model().from_bytes(model_bytes)
new_model.from_disk(model_path)