diff --git a/.jenkins/pytorch/test.sh b/.jenkins/pytorch/test.sh index 1e4a9a52066..a74a3e81580 100755 --- a/.jenkins/pytorch/test.sh +++ b/.jenkins/pytorch/test.sh @@ -450,6 +450,8 @@ test_xla() { # nightly version. test_forward_backward_compatibility() { set -x + # create a dummy ts model at this version + python test/create_dummy_torchscript_model.py /tmp/model_new.pt pushd test/forward_backward_compatibility python -m venv venv # shellcheck disable=SC1091 @@ -457,10 +459,21 @@ test_forward_backward_compatibility() { pip_install --pre torch -f https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html pip show torch python dump_all_function_schemas.py --filename nightly_schemas.txt + # FC: verify newmodel can be load with old code. + if ! python ../load_torchscript_model.py /tmp/model_new.pt; then + echo "FC check failed: new model cannot be load in old code" + return 1 + fi + python ../create_dummy_torchscript_model.py /tmp/model_old.pt deactivate rm -r venv pip show torch python check_forward_backward_compatibility.py --existing-schemas nightly_schemas.txt + # BC: verify old model can be load with new code + if ! python ../load_torchscript_model.py /tmp/model_old.pt; then + echo "BC check failed: old model cannot be load in new code" + return 1 + fi popd set +x assert_git_not_dirty diff --git a/test/create_dummy_torchscript_model.py b/test/create_dummy_torchscript_model.py new file mode 100644 index 00000000000..ffd869e27f0 --- /dev/null +++ b/test/create_dummy_torchscript_model.py @@ -0,0 +1,28 @@ +# Usage: python create_dummy_model.py +import sys +import torch +from torch import nn + + +class NeuralNetwork(nn.Module): + + def __init__(self): + super(NeuralNetwork, self).__init__() + self.flatten = nn.Flatten() + self.linear_relu_stack = nn.Sequential( + nn.Linear(28 * 28, 512), + nn.ReLU(), + nn.Linear(512, 512), + nn.ReLU(), + nn.Linear(512, 10), + ) + + def forward(self, x): + x = self.flatten(x) + logits = self.linear_relu_stack(x) + return logits + + +if __name__ == '__main__': + jit_module = torch.jit.script(NeuralNetwork()) + torch.jit.save(jit_module, sys.argv[1]) diff --git a/test/load_torchscript_model.py b/test/load_torchscript_model.py new file mode 100644 index 00000000000..dc8d4159d7f --- /dev/null +++ b/test/load_torchscript_model.py @@ -0,0 +1,6 @@ +import sys +import torch + +if __name__ == '__main__': + print(torch.jit.load(sys.argv[1])) + sys.exit(0)