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)
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)
def set_pytorch_transformer(model, transformer):
if model.attrs["has_transformer"]:
raise ValueError("Cannot set second transformer.")
model.layers.append(
PyTorchWrapper(
transformer,
convert_inputs=_convert_transformer_inputs,
convert_outputs=_convert_transformer_outputs,
))
model.attrs["has_transformer"] = True
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)
def create_wrapped_pytorch(width, dropout, nI, nO):
import torch
import torch.nn
import torch.nn.functional as F
class PyTorchModel(torch.nn.Module):
def __init__(self, width, nO, nI, dropout):
super(PyTorchModel, self).__init__()
self.dropout1 = torch.nn.Dropout2d(dropout)
self.dropout2 = torch.nn.Dropout2d(dropout)
self.fc1 = torch.nn.Linear(nI, width)
self.fc2 = torch.nn.Linear(width, nO)
def forward(self, x):
x = F.relu(x)
x = self.dropout1(x)
x = self.fc1(x)
x = F.relu(x)
x = self.dropout2(x)
x = self.fc2(x)
output = F.log_softmax(x, dim=1)
return output
return PyTorchWrapper(PyTorchModel(width, nO, nI, dropout))
def Transformer(name: str) -> Model[TokensPlus, List[Floats2d]]:
return PyTorchWrapper(
AutoModel.from_pretrained(name),
convert_inputs=convert_transformer_inputs,
convert_outputs=convert_transformer_outputs,
)