PyTorch 2.2 中文官方教程（二十）（1）

import torch
# use deeplabv3_resnet50 instead of deeplabv3_resnet101 to reduce the model size
model = torch.hub.load('pytorch/vision:v0.8.0', 'deeplabv3_resnet50', pretrained=True)
model.eval()
scriptedm = torch.jit.script(model)
torch.jit.save(scriptedm, "deeplabv3_scripted.pt") 

from PIL import Image
from torchvision import transforms
input_image = Image.open("deeplab.jpg")
preprocess = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
input_tensor = preprocess(input_image)
input_batch = input_tensor.unsqueeze(0)
with torch.no_grad():
    output = model(input_batch)['out'][0]
print(input_batch.shape)
print(output.shape) 

torch.Size([1, 3, 400, 400])
torch.Size([21, 400, 400]) 

class ViewController: UIViewController {
    var image = UIImage(named: "deeplab.jpg")!
    override func viewDidLoad() {
        super.viewDidLoad()
    }
    private lazy var module: TorchModule = {
        if let filePath = Bundle.main.path(forResource: "deeplabv3_scripted",
              ofType: "pt"),
            let module = TorchModule(fileAtPath: filePath) {
            return module
        } else {
            fatalError("Can't load the model file!")
        }
    }()
} 

- (unsigned  char*)predictImage:(void*)imageBuffer  {
  // 1\. the example deeplab.jpg size is size 400x400 and there are 21 semantic classes
  const  int  WIDTH  =  400;
  const  int  HEIGHT  =  400;
  const  int  CLASSNUM  =  21;
  at::Tensor  tensor  =  torch::from_blob(imageBuffer,  {1,  3,  WIDTH,  HEIGHT},  at::kFloat);
  torch::autograd::AutoGradMode  guard(false);
  at::AutoNonVariableTypeMode  non_var_type_mode(true);
  // 2\. convert the input tensor to an NSMutableArray for debugging
  float*  floatInput  =  tensor.data_ptr<float>();
  if  (!floatInput)  {
  return  nil;
  }
  NSMutableArray*  inputs  =  [[NSMutableArray  alloc]  init];
  for  (int  i  =  0;  i  <  3  *  WIDTH  *  HEIGHT;  i++)  {
  [inputs  addObject:@(floatInput[i])];
  }
  // 3\. the output of the model is a dictionary of string and tensor, as
  // specified at https://pytorch.org/hub/pytorch_vision_deeplabv3_resnet101
  auto  outputDict  =  _impl.forward({tensor}).toGenericDict();
  // 4\. convert the output to another NSMutableArray for easy debugging
  auto  outputTensor  =  outputDict.at("out").toTensor();
  float*  floatBuffer  =  outputTensor.data_ptr<float>();
  if  (!floatBuffer)  {
  return  nil;
  }
  NSMutableArray*  results  =  [[NSMutableArray  alloc]  init];
  for  (int  i  =  0;  i  <  CLASSNUM  *  WIDTH  *  HEIGHT;  i++)  {
  [results  addObject:@(floatBuffer[i])];
  }
  return  nil;
} 

// see the 20 semantic classes link in Introduction
const  int  DOG  =  12;
const  int  PERSON  =  15;
const  int  SHEEP  =  17;
NSMutableData*  data  =  [NSMutableData  dataWithLength:
  sizeof(unsigned  char)  *  3  *  WIDTH  *  HEIGHT];
unsigned  char*  buffer  =  (unsigned  char*)[data  mutableBytes];
// go through each element in the output of size [WIDTH, HEIGHT] and
// set different color for different classnum
for  (int  j  =  0;  j  <  WIDTH;  j++)  {
  for  (int  k  =  0;  k  <  HEIGHT;  k++)  {
  // maxi: the index of the 21 CLASSNUM with the max probability
  int  maxi  =  0,  maxj  =  0,  maxk  =  0;
  float  maxnum  =  -100000.0;
  for  (int  i  =  0;  i  <  CLASSNUM;  i++)  {
  if  ([results[i  *  (WIDTH  *  HEIGHT)  +  j  *  WIDTH  +  k]  floatValue]  >  maxnum)  {
  maxnum  =  [results[i  *  (WIDTH  *  HEIGHT)  +  j  *  WIDTH  +  k]  floatValue];
  maxi  =  i;  maxj  =  j;  maxk  =  k;
  }
  }
  int  n  =  3  *  (maxj  *  width  +  maxk);
  // color coding for person (red), dog (green), sheep (blue)
  // black color for background and other classes
  buffer[n]  =  0;  buffer[n+1]  =  0;  buffer[n+2]  =  0;
  if  (maxi  ==  PERSON)  buffer[n]  =  255;
  else  if  (maxi  ==  DOG)  buffer[n+1]  =  255;
  else  if  (maxi  ==  SHEEP)  buffer[n+2]  =  255;
  }
}
return  buffer;