Merge pull request #67 from jeffxtang/lite_updates

Demo app updates for Lite Interpreter: scripts, iOS code, and README

Author: xta0

D2Go/D2Go.xcodeproj/project.pbxproj

@@ -8,13 +8,13 @@
 
 /* Begin PBXBuildFile section */
 		260F86432627C445008F2567 /* classes.txt in Resources */ = {isa = PBXBuildFile; fileRef = 260F86422627C445008F2567 /* classes.txt */; };
+		26A8C10A26DFEC4700F4A58D /* d2go_optimized.ptl in Resources */ = {isa = PBXBuildFile; fileRef = 26A8C10926DFEC4700F4A58D /* d2go_optimized.ptl */; };
 		26C3345C26210C3E00C37AD3 /* AppDelegate.swift in Sources */ = {isa = PBXBuildFile; fileRef = 26C3345B26210C3E00C37AD3 /* AppDelegate.swift */; };
 		26C3345E26210C3E00C37AD3 /* SceneDelegate.swift in Sources */ = {isa = PBXBuildFile; fileRef = 26C3345D26210C3E00C37AD3 /* SceneDelegate.swift */; };
 		26C3346026210C3E00C37AD3 /* ViewController.swift in Sources */ = {isa = PBXBuildFile; fileRef = 26C3345F26210C3E00C37AD3 /* ViewController.swift */; };
 		26C3346326210C3E00C37AD3 /* Main.storyboard in Resources */ = {isa = PBXBuildFile; fileRef = 26C3346126210C3E00C37AD3 /* Main.storyboard */; };
 		26C3346526210C3F00C37AD3 /* Assets.xcassets in Resources */ = {isa = PBXBuildFile; fileRef = 26C3346426210C3F00C37AD3 /* Assets.xcassets */; };
 		26C3346826210C3F00C37AD3 /* LaunchScreen.storyboard in Resources */ = {isa = PBXBuildFile; fileRef = 26C3346626210C3F00C37AD3 /* LaunchScreen.storyboard */; };
-		26C3347626211FC900C37AD3 /* d2go_optimized.pt in Resources */ = {isa = PBXBuildFile; fileRef = 26C3347526211FC900C37AD3 /* d2go_optimized.pt */; };
 		26C3348C262120CB00C37AD3 /* PrePostProcessor.swift in Sources */ = {isa = PBXBuildFile; fileRef = 26C3347D262120CB00C37AD3 /* PrePostProcessor.swift */; };
 		26C3348D262120CB00C37AD3 /* UIImage+Helper.swift in Sources */ = {isa = PBXBuildFile; fileRef = 26C3347E262120CB00C37AD3 /* UIImage+Helper.swift */; };
 		26C3348F262120CB00C37AD3 /* ObjectDetector.swift in Sources */ = {isa = PBXBuildFile; fileRef = 26C33483262120CB00C37AD3 /* ObjectDetector.swift */; };
@@ -30,6 +30,7 @@
 
 /* Begin PBXFileReference section */
 		260F86422627C445008F2567 /* classes.txt */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = text; path = classes.txt; sourceTree = "<group>"; };
+		26A8C10926DFEC4700F4A58D /* d2go_optimized.ptl */ = {isa = PBXFileReference; lastKnownFileType = file; path = d2go_optimized.ptl; sourceTree = "<group>"; };
 		26C3345826210C3E00C37AD3 /* D2Go.app */ = {isa = PBXFileReference; explicitFileType = wrapper.application; includeInIndex = 0; path = D2Go.app; sourceTree = BUILT_PRODUCTS_DIR; };
 		26C3345B26210C3E00C37AD3 /* AppDelegate.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = AppDelegate.swift; sourceTree = "<group>"; };
 		26C3345D26210C3E00C37AD3 /* SceneDelegate.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = SceneDelegate.swift; sourceTree = "<group>"; };
@@ -38,7 +39,6 @@
 		26C3346426210C3F00C37AD3 /* Assets.xcassets */ = {isa = PBXFileReference; lastKnownFileType = folder.assetcatalog; path = Assets.xcassets; sourceTree = "<group>"; };
 		26C3346726210C3F00C37AD3 /* Base */ = {isa = PBXFileReference; lastKnownFileType = file.storyboard; name = Base; path = Base.lproj/LaunchScreen.storyboard; sourceTree = "<group>"; };
 		26C3346926210C3F00C37AD3 /* Info.plist */ = {isa = PBXFileReference; lastKnownFileType = text.plist.xml; path = Info.plist; sourceTree = "<group>"; };
-		26C3347526211FC900C37AD3 /* d2go_optimized.pt */ = {isa = PBXFileReference; lastKnownFileType = file; path = d2go_optimized.pt; sourceTree = "<group>"; };
 		26C3347D262120CB00C37AD3 /* PrePostProcessor.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = PrePostProcessor.swift; sourceTree = "<group>"; };
 		26C3347E262120CB00C37AD3 /* UIImage+Helper.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = "UIImage+Helper.swift"; sourceTree = "<group>"; };
 		26C33483262120CB00C37AD3 /* ObjectDetector.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = ObjectDetector.swift; sourceTree = "<group>"; };
@@ -105,7 +105,7 @@
 				26C3348A262120CB00C37AD3 /* test1.png */,
 				26C33484262120CB00C37AD3 /* test2.jpg */,
 				26C3348B262120CB00C37AD3 /* test3.png */,
-				26C3347526211FC900C37AD3 /* d2go_optimized.pt */,
+				26A8C10926DFEC4700F4A58D /* d2go_optimized.ptl */,
 				260F86422627C445008F2567 /* classes.txt */,
 			);
 			path = D2Go;
@@ -204,10 +204,10 @@
 				26C33490262120CB00C37AD3 /* test2.jpg in Resources */,
 				26C3346526210C3F00C37AD3 /* Assets.xcassets in Resources */,
 				26C3346326210C3E00C37AD3 /* Main.storyboard in Resources */,
-				26C3347626211FC900C37AD3 /* d2go_optimized.pt in Resources */,
 				26C33495262120CB00C37AD3 /* test1.png in Resources */,
 				26C33496262120CB00C37AD3 /* test3.png in Resources */,
 				260F86432627C445008F2567 /* classes.txt in Resources */,
+				26A8C10A26DFEC4700F4A58D /* d2go_optimized.ptl in Resources */,
 			);
 			runOnlyForDeploymentPostprocessing = 0;
 		};
@@ -401,9 +401,9 @@
 					"-l\"torch_cpu\"",
 					"-l\"torchvision_ops\"",
 					"-force_load",
-					"$(PODS_ROOT)/LibTorch/install/lib/libtorch.a",
+					"$(PODS_ROOT)/LibTorch-Lite/install/lib/libtorch.a",
 					"-force_load",
-					"$(PODS_ROOT)/LibTorch/install/lib/libtorch_cpu.a",
+					"$(PODS_ROOT)/LibTorch-Lite/install/lib/libtorch_cpu.a",
 					"-force_load",
 					"$(PODS_ROOT)/LibTorchvision/install/lib/libtorchvision_ops.a",
 				);
@@ -445,9 +445,9 @@
 					"-l\"torch_cpu\"",
 					"-l\"torchvision_ops\"",
 					"-force_load",
-					"$(PODS_ROOT)/LibTorch/install/lib/libtorch.a",
+					"$(PODS_ROOT)/LibTorch-Lite/install/lib/libtorch.a",
 					"-force_load",
-					"$(PODS_ROOT)/LibTorch/install/lib/libtorch_cpu.a",
+					"$(PODS_ROOT)/LibTorch-Lite/install/lib/libtorch_cpu.a",
 					"-force_load",
 					"$(PODS_ROOT)/LibTorchvision/install/lib/libtorchvision_ops.a",
 				);

D2Go/D2Go/Inference/InferenceModule.h

@@ -5,6 +5,7 @@
 // LICENSE file in the root directory of this source tree.
 
 #import <Foundation/Foundation.h>
+#import <UIKit/UIKit.h>
 
 NS_ASSUME_NONNULL_BEGIN
 

D2Go/D2Go/Inference/InferenceModule.mm

@@ -5,23 +5,22 @@
 // LICENSE file in the root directory of this source tree.
 
 #import "InferenceModule.h"
-#import <LibTorch.h>
+#import <Libtorch-Lite/Libtorch-Lite.h>
 
 const int input_width = 640;
 const int input_height = 640;
 const int threshold = 0.5;
 
 
 @implementation InferenceModule {
-    @protected torch::jit::script::Module _impl;
+    @protected torch::jit::mobile::Module _impl;
 }
 
 - (nullable instancetype)initWithFileAtPath:(NSString*)filePath {
     self = [super init];
     if (self) {
         try {
-            _impl = torch::jit::load(filePath.UTF8String);
-            _impl.eval();
+            _impl = torch::jit::_load_for_mobile(filePath.UTF8String);
         } catch (const std::exception& exception) {
             NSLog(@"%s", exception.what());
             return nil;
@@ -33,14 +32,18 @@ - (nullable instancetype)initWithFileAtPath:(NSString*)filePath {
 - (NSArray<NSNumber*>*)detectImage:(void*)imageBuffer {
     try {
         at::Tensor tensor = torch::from_blob(imageBuffer, { 3, input_width, input_height }, at::kFloat);
-        torch::autograd::AutoGradMode guard(false);
-        at::AutoNonVariableTypeMode non_var_type_mode(true);
-        
+        c10::InferenceMode guard;
+
         std::vector<torch::Tensor> v;
         v.push_back(tensor);
-        
+
+
+        CFTimeInterval startTime = CACurrentMediaTime();
         auto outputTuple = _impl.forward({ at::TensorList(v) }).toTuple();
-        
+        CFTimeInterval elapsedTime = CACurrentMediaTime() - startTime;
+        NSLog(@"inference time:%f", elapsedTime);
+
+
         auto outputDict = outputTuple->elements()[1].toList().get(0).toGenericDict();
         auto boxesTensor = outputDict.at("boxes").toTensor();
         auto scoresTensor = outputDict.at("scores").toTensor();
@@ -58,7 +61,7 @@ - (nullable instancetype)initWithFileAtPath:(NSString*)filePath {
         if (!labelsBuffer) {
             return nil;
         }
-        
+
         NSMutableArray* results = [[NSMutableArray alloc] init];
         long num = scoresTensor.numel();
         for (int i = 0; i < num; i++) {
@@ -72,9 +75,9 @@ - (nullable instancetype)initWithFileAtPath:(NSString*)filePath {
             [results addObject:@(scoresBuffer[i])];
             [results addObject:@(labelsBuffer[i])];
         }
-        
+
         return [results copy];
-        
+
     } catch (const std::exception& exception) {
         NSLog(@"%s", exception.what());
     }

D2Go/D2Go/Inference/ObjectDetector.swift

@@ -8,7 +8,7 @@ import UIKit
 
 class ObjectDetector {
     lazy var module: InferenceModule = {
-        if let filePath = Bundle.main.path(forResource: "d2go_optimized", ofType: "pt"),
+        if let filePath = Bundle.main.path(forResource: "d2go_optimized", ofType: "ptl"),
             let module = InferenceModule(fileAtPath: filePath) {
             return module
         } else {

D2Go/D2Go/d2go_optimized.pt

@@ -4,7 +4,7 @@ target 'D2Go' do
   # Comment the next line if you don't want to use dynamic frameworks
   use_frameworks!
 
-  pod 'LibTorch', '1.9.0'
+  pod 'LibTorch-Lite', '~>1.9.0'
   pod 'LibTorchvision', '0.10.0'
 
 end

D2Go/D2Go/d2go_optimized.ptl

@@ -2,22 +2,22 @@
 
 ## Introduction
 
-[Detectron2](https://github.com/facebookresearch/detectron2) is one of the most widely adopted open source projects and implements state-of-the-art object detection, semantic segmentation, panoptic segmentation, and human pose prediction. [D2Go](https://github.com/facebookresearch/d2go) is powered by PyTorch 1.9.0, torchvision 0.10.0, and Detectron2 with built-in SOTA networks for mobile - the D2Go model is very small (only 2.15MB) and runs very fast on iOS.
+[Detectron2](https://github.com/facebookresearch/detectron2) is one of the most widely adopted open source projects and implements state-of-the-art object detection, semantic segmentation, panoptic segmentation, and human pose prediction. [D2Go](https://github.com/facebookresearch/d2go) is powered by PyTorch 1.9, torchvision 0.10, and Detectron2 with built-in SOTA networks for mobile - the D2Go model is very small (only 2.15MB) and runs very fast on iOS.
 
 This D2Go iOS demo app shows how to prepare and use the D2Go model on iOS with the newly released LibTorchvision Cocoapods. The code is based on a previous PyTorch iOS [Object Detection demo app](https://github.com/pytorch/ios-demo-app/tree/master/ObjectDetection) that uses a pre-trained YOLOv5 model, with modified pre-processing and post-processing code required by the D2Go model.
 
 ## Prerequisites
 
 * PyTorch 1.9 and torchvision 0.10 (Optional)
 * Python 3.8 or above (Optional)
-* iOS Cocoapods LibTorch 1.9.0 and LibTorchvision 0.10.0
+* iOS Cocoapods LibTorch-Lite 1.9.0 and LibTorchvision 0.10.0
 * Xcode 12.4 or later
 
 ## Quick Start
 
 This section shows how to create and use the D2Go model in an iOS app. To just build and run the app without creating the D2Go model yourself, go directly to Step 4.
 
-1. Install PyTorch 1.9.0 and torchvision 0.10.0, for example:
+1. Install PyTorch 1.9 and torchvision 0.10, for example:
 
 ```
 conda create -n d2go python=3.8.5
@@ -50,6 +50,7 @@ Run the following command to create the optimized D2Go model `d2go_optimized.pt`
 ```
 python create_d2go.py
 ```
+Both the optimized JIT model and the Lite Interpreter model will be created and saved in the project folder.
 
 4. Build and run the D2Go iOS app
 

D2Go/Podfile

@@ -70,6 +70,7 @@ def forward(self, inputs: List[torch.Tensor]):
     scripted_model = torch.jit.script(wrapped_model)
     optimized_model = optimize_for_mobile(scripted_model)
     optimized_model.save("D2Go/d2go_optimized.pt")
+    optimized_model._save_for_lite_interpreter("D2Go/d2go_optimized.ptl")
 
 if __name__ == '__main__':
     test_export_torchvision_format()

D2Go/README.md

@@ -20,7 +20,6 @@ - (nullable instancetype)initWithFileAtPath:(NSString*)filePath {
     if (self) {
         try {
             _impl = torch::jit::_load_for_mobile(filePath.UTF8String);
-            _impl.eval();
         } catch (const std::exception& exception) {
             NSLog(@"%s", exception.what());
             return nil;
@@ -40,9 +39,6 @@ - (unsigned char*)segmentImage:(void *)imageBuffer withWidth:(int)width withHeig
 
         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);
-
         float* floatInput = tensor.data_ptr<float>();
         if (!floatInput) {
             return nil;
@@ -52,8 +48,13 @@ - (unsigned char*)segmentImage:(void *)imageBuffer withWidth:(int)width withHeig
             [inputs addObject:@(floatInput[i])];
         }
 
+        c10::InferenceMode guard;
+        
+        CFTimeInterval startTime = CACurrentMediaTime();
         auto outputDict = _impl.forward({ tensor }).toGenericDict();
-
+        CFTimeInterval elapsedTime = CACurrentMediaTime() - startTime;
+        NSLog(@"inference time:%f", elapsedTime);
+        
         auto outputTensor = outputDict.at("out").toTensor();
 
         float* floatBuffer = outputTensor.data_ptr<float>();