vector field example

Author: unixpickle

README.md

@@ -9,6 +9,10 @@ Problems solved by this project:
 
 See [Usage](#usage) for a high-level example of using the library. See [How it works](#how-it-works) to see which modules are supported.
 
+For fun, here's a vector field produced by differentiating the probability predictions of a two-class SVM (produced by [this script](examples/svm_vector_field.py)):
+
+<img src="examples/svm_vector_field.png" width="300" alt="A vector field quiver plot with two modes">
+
 # Usage
 
 First, train a model with scikit-learn as usual:
@@ -47,6 +51,8 @@ torch.jit.script(torch_model).save("path.pt")
 loaded_model = torch.jit.load("path.pt")
 ```
 
+For a full example of training a model and using its PyTorch translation, see [examples/svm_vector_field.py](examples/svm_vector_field.py).
+
 # How it works
 
 **sk2torch** contains PyTorch re-implementations of supported scikit-learn models. For a supported estimator `X`, a class `TorchX` in sk2torch will be able to read the attributes of `X` and convert them to `torch.Tensor` or simple Python types. `TorchX` subclasses `torch.nn.Module` and has a method for each inference API of `X` (e.g. `predict`, `decision_function`, etc.).

examples/svm_vector_field.png

@@ -0,0 +1,43 @@
+"""
+Train an SVM on a 2D classification problem, then plot a gradient vector field
+for the predicted class probabilty.
+"""
+
+import matplotlib.pyplot as plt
+import numpy as np
+import sk2torch
+import torch
+from sklearn.svm import SVC
+
+# Create a dataset of two Gaussians. There will be some overlap
+# between the two classes, which adds some uncertainty to the model.
+xs = np.concatenate(
+    [
+        np.random.random(size=(256, 2)) + [1, 0],
+        np.random.random(size=(256, 2)) + [-1, 0],
+    ],
+    axis=0,
+)
+ys = np.array([False] * 256 + [True] * 256)
+
+# Train an SVM on the data and wrap it in PyTorch.
+sk_model = SVC(probability=True)
+sk_model.fit(xs, ys)
+model = sk2torch.wrap(sk_model)
+
+# Create a coordinate grid to compute a vector field on.
+spaced = np.linspace(-2, 2, num=25)
+grid_xs = torch.tensor([[x, y] for x in spaced for y in spaced], requires_grad=True)
+
+# Compute the gradients of the SVM output.
+outputs = model.predict_proba(grid_xs)[:, 1]
+(input_grads,) = torch.autograd.grad(outputs.sum(), (grid_xs,))
+
+# Create a quiver plot of the vector field.
+plt.quiver(
+    grid_xs[:, 0].detach().numpy(),
+    grid_xs[:, 1].detach().numpy(),
+    input_grads[:, 0].detach().numpy(),
+    input_grads[:, 1].detach().numpy(),
+)
+plt.savefig("svm_vector_field.png")