chore: import upstream snapshot with attribution

examples/pytorch/torchscript/IrisClassification/MLproject

@@ -0,0 +1,12 @@
+name: iris-classification
+
+python_env: python_env.yaml
+
+entry_points:
+  main:
+    parameters:
+      epochs: {type: int, default: 50}
+
+    command: |
+          python iris_classification.py \
+            --epochs {epochs}

examples/pytorch/torchscript/IrisClassification/README.md

@@ -0,0 +1,43 @@
+## Iris classification example with MLflow
+
+This example demonstrates training a classification model on the Iris dataset, scripting the model with TorchScript, logging the
+scripted model to MLflow using
+[`mlflow.pytorch.log_model`](https://mlflow.org/docs/latest/python_api/mlflow.pytorch.html#mlflow.pytorch.log_model), and
+loading it back for inference using
+[`mlflow.pytorch.load_model`](https://mlflow.org/docs/latest/python_api/mlflow.pytorch.html#mlflow.pytorch.load_model)
+
+### Running the code
+
+To run the example via MLflow, navigate to the `mlflow/examples/pytorch/torchscript/IrisClassification` directory and run the command
+
+```
+mlflow run .
+```
+
+This will run `iris_classification.py` with the default set of parameters such as `--max_epochs=5`. You can see the default value in the `MLproject` file.
+
+In order to run the file with custom parameters, run the command
+
+```
+mlflow run . -P epochs=X
+```
+
+where `X` is your desired value for `epochs`.
+
+If you have the required modules for the file and would like to skip the creation of a conda environment, add the argument `--env-manager=local`.
+
+```
+mlflow run . --env-manager=local
+```
+
+Once the code is finished executing, you can view the run's metrics, parameters, and details by running the command
+
+```
+mlflow server
+```
+
+and navigating to [http://localhost:5000](http://localhost:5000).
+
+## Running against a custom tracking server
+
+To configure MLflow to log to a custom (non-default) tracking location, set the `MLFLOW_TRACKING_URI` environment variable, e.g. via `export MLFLOW_TRACKING_URI=http://localhost:5000/`. For more details, see [the docs](https://mlflow.org/docs/latest/tracking.html#where-runs-are-recorded)

examples/pytorch/torchscript/IrisClassification/iris_classification.py

@@ -0,0 +1,105 @@
+import argparse
+
+import torch
+import torch.nn.functional as F
+from sklearn.datasets import load_iris
+from sklearn.metrics import accuracy_score
+from sklearn.model_selection import train_test_split
+from torch import nn
+
+import mlflow.pytorch
+from mlflow.models import infer_signature
+
+
+class IrisClassifier(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.fc1 = nn.Linear(4, 10)
+        self.fc2 = nn.Linear(10, 10)
+        self.fc3 = nn.Linear(10, 3)
+
+    def forward(self, x):
+        x = F.relu(self.fc1(x))
+        x = F.relu(self.fc2(x))
+        x = F.dropout(x, 0.2)
+        x = self.fc3(x)
+        return x
+
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+def prepare_data():
+    iris = load_iris()
+    data = iris.data
+    labels = iris.target
+    target_names = iris.target_names
+
+    X_train, X_test, y_train, y_test = train_test_split(
+        data, labels, test_size=0.2, random_state=42, shuffle=True, stratify=labels
+    )
+
+    X_train = torch.FloatTensor(X_train).to(device)
+    X_test = torch.FloatTensor(X_test).to(device)
+    y_train = torch.LongTensor(y_train).to(device)
+    y_test = torch.LongTensor(y_test).to(device)
+
+    return X_train, X_test, y_train, y_test, target_names
+
+
+def train_model(model, epochs, X_train, y_train):
+    criterion = nn.CrossEntropyLoss()
+    optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
+
+    for epoch in range(epochs):
+        out = model(X_train)
+        loss = criterion(out, y_train).to(device)
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+
+        if epoch % 10 == 0:
+            print("number of epoch", epoch, "loss", float(loss))
+
+    return model
+
+
+def test_model(model, X_test, y_test):
+    model.eval()
+    with torch.no_grad():
+        predict_out = model(X_test)
+        _, predict_y = torch.max(predict_out, 1)
+
+        print("\nprediction accuracy", float(accuracy_score(y_test.cpu(), predict_y.cpu())))
+        return infer_signature(X_test.numpy(), predict_out.numpy())
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Iris Classification Torchscripted model")
+
+    parser.add_argument(
+        "--epochs", type=int, default=100, help="number of epochs to run (default: 100)"
+    )
+
+    args = parser.parse_args()
+
+    model = IrisClassifier()
+    model = model.to(device)
+    X_train, X_test, y_train, y_test, target_names = prepare_data()
+    scripted_model = torch.jit.script(model)  # scripting the model
+    scripted_model = train_model(scripted_model, args.epochs, X_train, y_train)
+    signature = test_model(scripted_model, X_test, y_test)
+
+    with mlflow.start_run() as run:
+        mlflow.pytorch.log_model(
+            scripted_model, name="model", signature=signature
+        )  # logging scripted model
+        model_path = mlflow.get_artifact_uri("model")
+        loaded_pytorch_model = mlflow.pytorch.load_model(model_path)  # loading scripted model
+        model.eval()
+        with torch.no_grad():
+            test_datapoint = torch.Tensor([4.4000, 3.0000, 1.3000, 0.2000]).to(device)
+            prediction = loaded_pytorch_model(test_datapoint)
+            actual = "setosa"
+            predicted = target_names[torch.argmax(prediction)]
+            print(f"\nPREDICTION RESULT: ACTUAL: {actual}, PREDICTED: {predicted}")

examples/pytorch/torchscript/IrisClassification/python_env.yaml

@@ -0,0 +1,8 @@
+build_dependencies:
+  - pip
+dependencies:
+  - scikit-learn
+  - cloudpickle==1.6.0
+  - boto3
+  - torchvision>=0.9.1
+  - torch>=1.9.0

examples/pytorch/torchscript/MNIST/MLproject

@@ -0,0 +1,16 @@
+name: mnist-torchscript
+
+python_env: python_env.yaml
+
+entry_points:
+  main:
+    parameters:
+      epochs: {type: int, default: 5}
+      batch_size: {type: int, default: 64}
+      learning_rate: {type: float, default: 1e-3}
+
+    command: |
+          python mnist_torchscript.py \
+            --epochs {epochs} \
+            --batch-size {batch_size} \
+            --lr {learning_rate}

examples/pytorch/torchscript/MNIST/README.md

@@ -0,0 +1,50 @@
+## MNIST example with MLflow
+
+This example demonstrates training of MNIST handwritten recognition model and logging it as torch scripted model.
+`mlflow.pytorch.log_model()` is used to log the scripted model to MLflow and `mlflow.pytorch.load_model()` to load it from MLflow
+
+### Code related to MLflow:
+
+This will log the TorchScripted model into MLflow and load the logged model.
+
+## Setting Tracking URI
+
+MLflow tracking URI can be set using the environment variable `MLFLOW_TRACKING_URI`
+
+Example: `export MLFLOW_TRACKING_URI=http://localhost:5000/`
+
+For more details - https://mlflow.org/docs/latest/tracking.html#where-runs-are-recorded
+
+### Running the code
+
+To run the example via MLflow, navigate to the `mlflow/examples/pytorch/torchscript/MNIST` directory and run the command
+
+```
+mlflow run .
+```
+
+This will run `mnist_torchscript.py` with the default set of parameters such as `--max_epochs=5`. You can see the default value in the `MLproject` file.
+
+In order to run the file with custom parameters, run the command
+
+```
+mlflow run . -P epochs=X
+```
+
+where `X` is your desired value for `epochs`.
+
+If you have the required modules for the file and would like to skip the creation of a conda environment, add the argument `--env-manager=local`.
+
+```
+mlflow run . --env-manager=local
+```
+
+Once the code is finished executing, you can view the run's metrics, parameters, and details by running the command
+
+```
+mlflow server
+```
+
+and navigating to [http://localhost:5000](http://localhost:5000).
+
+For more information on MLflow tracking, click [here](https://www.mlflow.org/docs/latest/tracking.html#mlflow-tracking) to view documentation.

examples/pytorch/torchscript/MNIST/mnist_torchscript.py

@@ -0,0 +1,196 @@
+import argparse
+
+import torch
+import torch.nn.functional as F
+from torch import nn, optim
+from torch.optim.lr_scheduler import StepLR
+from torchvision import datasets, transforms
+
+import mlflow
+import mlflow.pytorch
+
+
+class Net(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv1 = nn.Conv2d(1, 32, 3, 1)
+        self.conv2 = nn.Conv2d(32, 64, 3, 1)
+        self.dropout1 = nn.Dropout(0.25)
+        self.dropout2 = nn.Dropout(0.5)
+        self.fc1 = nn.Linear(9216, 128)
+        self.fc2 = nn.Linear(128, 10)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = F.relu(x)
+        x = self.conv2(x)
+        x = F.relu(x)
+        x = F.max_pool2d(x, 2)
+        x = self.dropout1(x)
+        x = torch.flatten(x, 1)
+        x = self.fc1(x)
+        x = F.relu(x)
+        x = self.dropout2(x)
+        x = self.fc2(x)
+        x = F.log_softmax(x, dim=1)
+        return x
+
+
+def train(args, model, device, train_loader, optimizer, epoch):
+    model.train()
+    for batch_idx, (data, target) in enumerate(train_loader):
+        data = data.to(device)
+        target = target.to(device)
+        optimizer.zero_grad()
+        output = model(data)
+        loss = F.nll_loss(output, target)
+        loss.backward()
+        optimizer.step()
+        if batch_idx % args.log_interval == 0:
+            print(
+                "Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}".format(
+                    epoch,
+                    batch_idx * len(data),
+                    len(train_loader.dataset),
+                    100.0 * batch_idx / len(train_loader),
+                    loss.item(),
+                )
+            )
+            if args.dry_run:
+                break
+
+
+def test(model, device, test_loader):
+    model.eval()
+    test_loss = 0
+    correct = 0
+    with torch.no_grad():
+        for data, target in test_loader:
+            data = data.to(device)
+            target = target.to(device)
+            output = model(data)
+            test_loss += F.nll_loss(output, target, reduction="sum").item()  # sum up batch loss
+            pred = output.argmax(dim=1, keepdim=True)  # get the index of the max log-probability
+            correct += pred.eq(target.view_as(pred)).sum().item()
+
+    test_loss /= len(test_loader.dataset)
+
+    print(
+        "\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n".format(
+            test_loss,
+            correct,
+            len(test_loader.dataset),
+            100.0 * correct / len(test_loader.dataset),
+        )
+    )
+
+
+def main():
+    # Training settings
+    parser = argparse.ArgumentParser(description="PyTorch MNIST Example")
+    parser.add_argument(
+        "--batch-size",
+        type=int,
+        default=64,
+        metavar="N",
+        help="input batch size for training (default: 64)",
+    )
+    parser.add_argument(
+        "--test-batch-size",
+        type=int,
+        default=1000,
+        metavar="N",
+        help="input batch size for testing (default: 1000)",
+    )
+    parser.add_argument(
+        "--epochs",
+        type=int,
+        default=14,
+        metavar="N",
+        help="number of epochs to train (default: 14)",
+    )
+    parser.add_argument(
+        "--lr",
+        type=float,
+        default=1.0,
+        metavar="LR",
+        help="learning rate (default: 1.0)",
+    )
+    parser.add_argument(
+        "--gamma",
+        type=float,
+        default=0.7,
+        metavar="M",
+        help="Learning rate step gamma (default: 0.7)",
+    )
+    parser.add_argument(
+        "--no-cuda", action="store_true", default=False, help="disables CUDA training"
+    )
+    parser.add_argument(
+        "--dry-run",
+        action="store_true",
+        default=False,
+        help="quickly check a single pass",
+    )
+    parser.add_argument("--seed", type=int, default=1, metavar="S", help="random seed (default: 1)")
+    parser.add_argument(
+        "--log-interval",
+        type=int,
+        default=10,
+        metavar="N",
+        help="how many batches to wait before logging training status",
+    )
+    parser.add_argument(
+        "--save-model",
+        action="store_true",
+        default=False,
+        help="For Saving the current model",
+    )
+
+    args = parser.parse_args()
+    use_cuda = not args.no_cuda and torch.cuda.is_available()
+
+    torch.manual_seed(args.seed)
+
+    device = torch.device("cuda" if use_cuda else "cpu")
+
+    train_kwargs = {"batch_size": args.batch_size}
+    test_kwargs = {"batch_size": args.test_batch_size}
+    if use_cuda:
+        cuda_kwargs = {"num_workers": 1, "pin_memory": True, "shuffle": True}
+        train_kwargs.update(cuda_kwargs)
+        test_kwargs.update(cuda_kwargs)
+
+    transform = transforms.Compose([
+        transforms.ToTensor(),
+        transforms.Normalize((0.1307,), (0.3081,)),
+    ])
+    dataset1 = datasets.MNIST("../data", train=True, download=True, transform=transform)
+    dataset2 = datasets.MNIST("../data", train=False, transform=transform)
+    train_loader = torch.utils.data.DataLoader(dataset1, **train_kwargs)
+    test_loader = torch.utils.data.DataLoader(dataset2, **test_kwargs)
+
+    model = Net().to(device)
+    scripted_model = torch.jit.script(model)  # scripting the model
+    optimizer = optim.Adadelta(model.parameters(), lr=args.lr)
+
+    scheduler = StepLR(optimizer, step_size=1, gamma=args.gamma)
+    for epoch in range(1, args.epochs + 1):
+        train(args, scripted_model, device, train_loader, optimizer, epoch)
+        scheduler.step()
+    test(scripted_model, device, test_loader)
+    with mlflow.start_run():
+        mlflow.pytorch.log_model(scripted_model, name="model")  # logging scripted model
+        model_path = mlflow.get_artifact_uri("model")
+        loaded_pytorch_model = mlflow.pytorch.load_model(model_path)  # loading scripted model
+        model.eval()
+        with torch.no_grad():
+            test_datapoint, test_target = next(iter(test_loader))
+            prediction = loaded_pytorch_model(test_datapoint[0].unsqueeze(0).to(device))
+            actual = test_target[0].item()
+            predicted = torch.argmax(prediction).item()
+            print(f"\nPREDICTION RESULT: ACTUAL: {actual!s}, PREDICTED: {predicted!s}")
+
+
+if __name__ == "__main__":
+    main()

examples/pytorch/torchscript/MNIST/python_env.yaml

@@ -0,0 +1,8 @@
+build_dependencies:
+  - pip
+dependencies:
+  - mlflow
+  - cloudpickle==1.6.0
+  - boto3
+  - torchvision>=0.9.1
+  - torch>=1.9.0