import torch
  class MyModule(torch.nn.Module):

    def __init__(self, N, M):
      super(MyModule, self).__init__()
      self.weight = torch.nn.Parameter(torch.rand(N, M))

    def forward(self, input):
      if input.sum() > 0:
        output = self.weight.mv(input)
      else:
        output = self.weight + input
      return output

    # Compile the model code to a static representation
    my_script_module = torch.jit.script(MyModule(3, 4))

    # Save the compiled code and model data so it can be loaded elsewhere
    my_script_module.save("my_script_module.pt")

Production Ready

With TorchScript, PyTorch provides ease-of-use and flexibility in eager mode, while seamlessly transitioning to graph mode for speed, optimization, and functionality in C++ runtime environments.

TorchServe

TorchServe is an easy to use tool for deploying PyTorch models at scale. It is cloud and environment agnostic and supports features such as multi-model serving, logging, metrics and the creation of RESTful endpoints for application integration.

  ## Convert the model from PyTorch to TorchServe format
  torch-model-archiver --model-name densenet161 \
  --version 1.0 --model-file serve/examples/image_classifier/densenet_161/model.py \
  --serialized-file densenet161-8d451a50.pth \
  --extra-files serve/examples/image_classifier/index_to_name.json \
  --handler image_classifier

  ## Host your PyTorch model

 torchserve --start --model-store model_store --models densenet161=densenet161.mar

  import torch.distributed as dist
  from torch.nn.parallel import DistributedDataParallel
  
  dist.init_process_group(backend='gloo')
  model = DistributedDataParallel(model)

Distributed Training

Optimize performance in both research and production by taking advantage of native support for asynchronous execution of collective operations and peer-to-peer communication that is accessible from Python and C++.

Mobile (Experimental)

PyTorch supports an end-to-end workflow from Python to deployment on iOS and Android. It extends the PyTorch API to cover common preprocessing and integration tasks needed for incorporating ML in mobile applications.

  ## Save your model
  torch.jit.script(model).save("my_mobile_model.pt")

  ## iOS prebuilt binary
  pod ‘LibTorch’
  ## Android prebuilt binary
  implementation 'org.pytorch:pytorch_android:1.3.0'

  ## Run your model (Android example)
  Tensor input = Tensor.fromBlob(data, new long[]{1, data.length});
  IValue output = module.forward(IValue.tensor(input));
  float[] scores = output.getTensor().getDataAsFloatArray();

  import torchvision.models as models
  resnet18 = models.resnet18(pretrained=True)
  alexnet = models.alexnet(pretrained=True)
  squeezenet = models.squeezenet1_0(pretrained=True)
  vgg16 = models.vgg16(pretrained=True)
  densenet = models.densenet161(pretrained=True)
  inception = models.inception_v3(pretrained=True)

Robust Ecosystem

An active community of researchers and developers have built a rich ecosystem of tools and libraries for extending PyTorch and supporting development in areas from computer vision to reinforcement learning.

Native ONNX Support

Export models in the standard ONNX (Open Neural Network Exchange) format for direct access to ONNX-compatible platforms, runtimes, visualizers, and more.

  import torch.onnx
  import torchvision

  dummy_input = torch.randn(1, 3, 224, 224)
  model = torchvision.models.alexnet(pretrained=True)
  torch.onnx.export(model, dummy_input, "alexnet.onnx")

  #include <torch/torch.h>

  torch::nn::Linear model(num_features, 1);
  torch::optim::SGD optimizer(model->parameters());
  auto data_loader = torch::data::data_loader(dataset);

  for (size_t epoch = 0; epoch < 10; ++epoch) {
    for (auto batch : data_loader) {
      auto prediction = model->forward(batch.data);
      auto loss = loss_function(prediction, batch.target);
      loss.backward();
      optimizer.step();
    }
  }

C++ Front-End

The C++ frontend is a pure C++ interface to PyTorch that follows the design and architecture of the established Python frontend. It is intended to enable research in high performance, low latency and bare metal C++ applications.

Cloud Support

PyTorch is well supported on major cloud platforms, providing frictionless development and easy scaling through prebuilt images, large scale training on GPUs, ability to run models in a production scale environment, and more.

  export IMAGE_FAMILY="pytorch-latest-cpu"
  export ZONE="us-west1-b"
  export INSTANCE_NAME="my-instance"
  
  gcloud compute instances create $INSTANCE_NAME \
    --zone=$ZONE \
    --image-family=$IMAGE_FAMILY \
    --image-project=deeplearning-platform-release