Core ML 和 Vision 教程：在 iOS 设备上直接进行模型微调

import CoreML
import Vision

// 定义了一个懒加载的图像分析请求
private lazy var classificationRequest: VNCoreMLRequest = {
  do {
    // 创建模型实例
    let model = try VNCoreMLModel(for: SqueezeNet(configuration: MLModelConfiguration()).model)
    // 基于该模型实例化一个图像分析请求对象。完成处理程序后，接收分类结果并打印出来。
    let request = VNCoreMLRequest(model: model) { request, _ in
        if let classifications =
          request.results as? [VNClassificationObservation] {
          print("Classification results: \(classifications)")
        }
    }
    // 使用Vision来裁剪输入的图像以符合模型的期望
    request.imageCropAndScaleOption = .centerCrop
    return request
  } catch {
    // 通过杀死应用程序来处理模型加载错误。模型是应用程序捆绑的一部分，应该不会走到这。实际应用中，埋个点吧。
    fatalError("Failed to load Vision ML model: \(error)")
  }
}()

func classifyImage(_ image: UIImage) {
  // 获取图像的方向，并且将图像转换成CIImage格式提供给后续步骤
  guard let orientation = CGImagePropertyOrientation(rawValue: UInt32(image.imageOrientation.rawValue)) else {
    return
  }
  guard let ciImage = CIImage(image: image) else {
    fatalError("Unable to create \(CIImage.self) from \(image).")
  }
  // 在一个后台队列中启动一个异步分类请求，当句柄在外部被创建并且安排上时，这个Vision请求将被执行
  DispatchQueue.global(qos: .userInitiated).async {
    let handler = VNImageRequestHandler(ciImage: ciImage, orientation: orientation)
    do {
      try handler.perform([self.classificationRequest])
    } catch {
      print("Failed to perform classification.\n\(error.localizedDescription)")
    }
  }
}

classifyImage(image)    // 当用户选择一张图片时，触发分类请求

// print("Classification results: \(classifications)")
let topClassifications = classifications.prefix(2).map {
  (confidence: $0.confidence, identifier: $0.identifier)
}
print("Top classifications: \(topClassifications)")

if let quote = getQuote() {
  quoteTextView.text = quote.text
}

func processClassifications(for request: VNRequest, error: Error?) {
  DispatchQueue.main.async {
    // 处理来自图像分类请求的结果
    if let classifications =
      request.results as? [VNClassificationObservation] {
      // 提取前两个预测结果
      let topClassifications = classifications.prefix(2).map {
        (confidence: $0.confidence, identifier: $0.identifier)
      }
      print("Top classifications: \(topClassifications)")
      let topIdentifiers =
        topClassifications.map {$0.identifier.lowercased() }
      // 将预测结果传入getQuote(for:)获得一个相关的箴言
      if let quote = self.getQuote(for: topIdentifiers) {
        self.quoteTextView.text = quote.text
      }
    }
  }
}

let request = VNCoreMLRequest(model: model) { [weak self] request, error in
    guard let self = self else {
      return
  }
  // 当预测结束时，调用方法更新箴言
  self.processClassifications(for: request, error: error)
}

performSegue(withIdentifier: "AddShortcutSegue", sender: self)

print("Training data ready for label: \(selectedEmoji ?? "")")
// 当用户点击保存按钮时，回到首页
performSegue(
  withIdentifier: "AddShortcutUnwindSegue",
  sender: self)

stickerLabel.isUserInteractionEnabled = true
let panGestureRecognizer = UIPanGestureRecognizer(
  target: self,
  action: #selector(handlePanGesture(_:)))
stickerLabel.addGestureRecognizer(panGestureRecognizer)

if #available(iOS 14.0, *) {
  canvasView.drawingPolicy = .anyInput
} else {
  canvasView.allowsFingerDrawing = true
}

var drawingView: DrawingView!   // 用户绘制快捷方式的画布视图

// 创建绘图视图示例
drawingView = DrawingView(frame: stickerView.bounds)
// 添加到主视图
view.addSubview(drawingView)
// 添加约束防止和贴纸视图重叠
drawingView.translatesAutoresizingMaskIntoConstraints = false
NSLayoutConstraint.activate([
  drawingView.topAnchor.constraint(equalTo: stickerView.topAnchor),
  drawingView.leftAnchor.constraint(equalTo: stickerView.leftAnchor),
  drawingView.rightAnchor.constraint(equalTo: stickerView.rightAnchor),
  drawingView.bottomAnchor.constraint(equalTo: stickerView.bottomAnchor)
])

addCanvasForDrawing()
drawingView.isHidden = true
// 添加绘图视图，初始隐藏

// 清除画布，隐藏绘图视图以便表情符号的贴纸可以正确展示
drawingView.clearCanvas()
drawingView.isHidden = false

import CoreML   // 引入机器学习框架

extension UpdatableDrawingClassifier {
  // 确保图像与模型所期望的一致
  var imageConstraint: MLImageConstraint {
    return model.modelDescription
      .inputDescriptionsByName["drawing"]!
      .imageConstraint!
  }

  // 用绘图的CVPixelBuffer调用模型的预测方法。
  // 返回预测的表情符号标签，如果没有匹配的则返回nil
  func predictLabelFor(_ value: MLFeatureValue) -> String? {
    guard
      let pixelBuffer = value.imageBufferValue,
      let prediction = try? prediction(drawing: pixelBuffer).label
      else {
         return nil
    }
    if prediction == "unknown" {
      print("No prediction found")
      return nil
    }
    return prediction
  }
}

// 映射可更新的模型
struct UpdatableModel {
  private static var updatedDrawingClassifier: UpdatableDrawingClassifier?
  private static let appDirectory = FileManager.default.urls(
    for: .applicationSupportDirectory,
    in: .userDomainMask).first!
  // 指向原始编译模型
  private static let defaultModelURL =
  UpdatableDrawingClassifier.urlOfModelInThisBundle
  // 保存模型的位置
  private static var updatedModelURL =
  appDirectory.appendingPathComponent("personalized.mlmodelc")
  private static var tempUpdatedModelURL =
  appDirectory.appendingPathComponent("personalized_tmp.mlmodelc")

  private init() { }

  static var imageConstraint: MLImageConstraint {
    guard let model = try? updatedDrawingClassifier ?? UpdatableDrawingClassifier(configuration: MLModelConfiguration()) else {
      fatalError("init UpdatableDrawingClassifier error")
    }
    return model.imageConstraint
  }
}

private extension UpdatableModel {
  // 加载模型
  static func loadModel() {
    let fileManager = FileManager.default
    if !fileManager.fileExists(atPath: updatedModelURL.path) {
      do {
        let updatedModelParentURL =
          updatedModelURL.deletingLastPathComponent()
        try fileManager.createDirectory(
          at: updatedModelParentURL,
          withIntermediateDirectories: true,
          attributes: nil)
        let toTemp = updatedModelParentURL
          .appendingPathComponent(defaultModelURL.lastPathComponent)
        try fileManager.copyItem(
          at: defaultModelURL,
          to: toTemp)
        try fileManager.moveItem(
          at: toTemp,
          to: updatedModelURL)
      } catch {
        print("Error: \(error)")
        return
      }
    }
    guard let model = try? UpdatableDrawingClassifier(
      contentsOf: updatedModelURL) else {
      return
    }
    // 模型加载到内存
    updatedDrawingClassifier = model
  }
}

extension UpdatableModel {
  static func predictLabelFor(_ value: MLFeatureValue) -> String? {
    loadModel()
    return updatedDrawingClassifier?.predictLabelFor(value)
  }
}

import CoreML

var featureValue: MLFeatureValue {
  let imageConstraint = UpdatableModel.imageConstraint

  let preparedImage = whiteTintedImage

  let imageFeatureValue =
    try? MLFeatureValue(cgImage: preparedImage, constraint: imageConstraint)
  return imageFeatureValue!
}

extension CreateQuoteViewController: DrawingViewDelegate {
  func drawingDidChange(_ drawingView: DrawingView) {
    // 绘图的边界，防止越界
    let drawingRect = drawingView.boundingSquare()
    // 绘图实例
    let drawing = Drawing(
      drawing: drawingView.canvasView.drawing,
      rect: drawingRect)
    // 为绘图预测输入创建特征值
    let imageFeatureValue = drawing.featureValue
    // 进行预测，以获得与该绘制图形相对应的表情符号
    let drawingLabel =
      UpdatableModel.predictLabelFor(imageFeatureValue)
    // 更新主队列中的视图，清除画布并将预测的表情符号添加到主视图中
    DispatchQueue.main.async {
      drawingView.clearCanvas()
      guard let emoji = drawingLabel else {
        return
      }
      self.addStickerToCanvas(emoji, at: drawingRect)
    }
  }
}

drawingView.clearCanvas()

drawingView.delegate = self

import CoreML

func prepareTrainingData() throws -> MLBatchProvider {
  // 初始化一个空的 MLFeatureProvider 数组
  var featureProviders: [MLFeatureProvider] = []
  // 定义模型训练输入的名称
  let inputName = "drawing"
  let outputName = "label"
  // 循环浏览数据存储中的图形数据
  for drawing in drawings {
    if let drawing = drawing {
      // 将绘图训练输入包在一个特征值中
      let inputValue = drawing.featureValue
      // 将emoji训练输入包在一个特征值中
      let outputValue = MLFeatureValue(string: emoji)
      // 为训练输入创建一个MLFeatureValue集合。这是一个训练输入名称和特征值的字典。
      let dataPointFeatures: [String: MLFeatureValue] =
        [inputName: inputValue,
        outputName: outputValue]
      // 为 MLFeatureValue 集合创建一个 MLFeatureProvider，并将其追加到数组中
      if let provider =
        try? MLDictionaryFeatureProvider(
          dictionary: dataPointFeatures) {
        featureProviders.append(provider)
      }
    }
  }
  // 最后，从MLFeatureProvider数组中创建一个批处理对象（MLArrayBatchProvider）
  return MLArrayBatchProvider(array: featureProviders)
}

static func updateModel(
  at url: URL,
  with trainingData: MLBatchProvider,
  completionHandler: @escaping (MLUpdateContext) -> Void
) {
  do {
    let updateTask = try MLUpdateTask(
      forModelAt: url,
      trainingData: trainingData,
      configuration: nil,
      completionHandler: completionHandler)
    updateTask.resume()
  } catch {
    print("Couldn't create an MLUpdateTask.")
  }
}

static func saveUpdatedModel(_ updateContext: MLUpdateContext) {
  // 首先，从内存中获取更新的模型。这与原始模型不一样。
  let updatedModel = updateContext.model
  let fileManager = FileManager.default
  do {
    // 然后，创建一个中间文件夹来保存更新的模型。
    try fileManager.createDirectory(
        at: tempUpdatedModelURL,
        withIntermediateDirectories: true,
        attributes: nil)
    // 把更新的模型写到一个临时文件夹中
    try updatedModel.write(to: tempUpdatedModelURL)
    // 替换模型文件夹的内容
    // 直接覆盖现有的mlmodelc文件夹会出现错误。
    // 解决方案是保存到一个中间文件夹，然后把内容复制过来。
    _ = try fileManager.replaceItemAt(
      updatedModelURL,
      withItemAt: tempUpdatedModelURL)
    print("Updated model saved to:\n\t\(updatedModelURL)")
  } catch let error {
    print("Could not save updated model to the file system: \(error)")
    return
  }
}

static func updateWith(
  trainingData: MLBatchProvider,
  completionHandler: @escaping () -> Void
) {
  loadModel()
  UpdatableDrawingClassifier.updateModel(
    at: updatedModelURL,
    with: trainingData) { context in
      saveUpdatedModel(context)
      DispatchQueue.main.async { completionHandler() }
  }
}

do {
  let trainingData = try drawingDataStore.prepareTrainingData()
  DispatchQueue.global(qos: .userInitiated).async {
    UpdatableModel.updateWith(trainingData: trainingData) {
      DispatchQueue.main.async {
        self.performSegue(
          withIdentifier: "AddShortcutUnwindSegue",
          sender: self)
      }
    }
  }
} catch {
  print("Error updating model", error)
}