--- /dev/null
+#include <torch/csrc/jit/backends/backend.h>
+#include <torch/csrc/jit/backends/coreml/objc/PTMCoreMLExecutor.h>
+#include <torch/script.h>
+
+#import <CoreML/CoreML.h>
+
+namespace torch {
+namespace jit {
+namespace mobile {
+namespace coreml {
+
+static constexpr int SUPPORTED_COREML_VER = 4;
+
+enum TensorType {
+ Float,
+ Double,
+ Int,
+ Long,
+ Undefined,
+};
+
+static inline c10::ScalarType scalarType(TensorType type) {
+ switch (type) {
+ case TensorType::Float:
+ return c10::ScalarType::Float;
+ case TensorType::Double:
+ return c10::ScalarType::Double;
+ case TensorType::Int:
+ return c10::ScalarType::Int;
+ case TensorType::Long:
+ return c10::ScalarType::Long;
+ case TensorType::Undefined:
+ return c10::ScalarType::Undefined;
+ default:
+ return c10::ScalarType::Undefined;
+ }
+}
+
+static id parse(NSString* jsonStr) {
+ NSData* data = [jsonStr dataUsingEncoding:NSUTF8StringEncoding];
+ NSError* error;
+ id result = [NSJSONSerialization JSONObjectWithData:data
+ options:0
+ error:&error];
+ if (error || !result) {
+ TORCH_CHECK(
+ false,
+ "parsing JSON string failed!",
+ error.localizedDescription.UTF8String);
+ }
+
+ return result;
+}
+
+struct TensorSpec {
+ public:
+ TensorSpec() = delete;
+ TensorSpec(NSArray<NSString*>* spec) {
+ TORCH_CHECK(spec.count == 3);
+ name_ = spec[0];
+ dtype_ = (TensorType)spec[1].intValue;
+ NSArray* sizes = parse(spec[2]);
+ for (NSString* dim in sizes) {
+ sizes_.emplace_back(dim.integerValue);
+ }
+ }
+ int64_t numel() const {
+ return std::accumulate(
+ begin(sizes_), end(sizes_), 1, std::multiplies<int64_t>());
+ }
+ NSString* name() {
+ return name_;
+ }
+ std::vector<int64_t> sizes() {
+ return sizes_;
+ }
+ TensorType dtype() {
+ return dtype_;
+ }
+
+ private:
+ NSString* name_ = @"";
+ TensorType dtype_ = TensorType::Float;
+ std::vector<int64_t> sizes_{};
+};
+
+struct CoreMLConfig {
+ public:
+ CoreMLConfig() = delete;
+ CoreMLConfig(NSDictionary* dict)
+ : coreMLVersion_([dict[@"spec_ver"] intValue]),
+ backend_([dict[@"backend"] lowercaseString]),
+ allow_low_precision_([dict[@"allow_low_precision"] boolValue]) {
+ TORCH_CHECK(
+ coreMLVersion_ >= SUPPORTED_COREML_VER,
+ "Only Core ML version 4 and above are supported");
+ }
+ int64_t coreMLVersion() const {
+ return coreMLVersion_;
+ }
+ NSString* backend() const {
+ return backend_;
+ }
+ bool allowLowPrecision() const {
+ return allow_low_precision_;
+ }
+
+ private:
+ int64_t coreMLVersion_ = SUPPORTED_COREML_VER;
+ NSString* backend_ = @"CPU";
+ bool allow_low_precision_ = true;
+};
+
+struct MetaData {
+ public:
+ MetaData(NSDictionary* dict)
+ : torchVer_(dict[@"torch_ver"]),
+ coremltoolVer_(dict[@"coremltool_ver"]) {}
+ NSString* torchVer() const {
+ return torchVer_;
+ }
+ NSString* coremltoolVer() const {
+ return coremltoolVer_;
+ }
+
+ private:
+ NSString* torchVer_ = @"";
+ NSString* coremltoolVer_ = @"";
+};
+
+// Wrap the Objective-C executor into a C++ to be able to pack into IValue
+struct API_AVAILABLE(ios(11.0), macos(10.13)) CoreMLExecutorWrapper
+ : public CustomClassHolder {
+ public:
+ CoreMLExecutorWrapper(
+ PTMCoreMLExecutor* executor,
+ std::vector<TensorSpec>& inputs,
+ std::vector<TensorSpec>& outputs,
+ CoreMLConfig config)
+ : executor_(executor),
+ inputs_(inputs),
+ outputs_(outputs),
+ config_(config) {}
+ c10::List<torch::Tensor> execute(c10::impl::GenericList inputs) {
+ std::vector<PTMCoreMLFeatureSpecs> inputSpecs;
+ std::vector<PTMCoreMLFeatureSpecs> outputSpecs;
+ int inputSpecIndex = 0;
+ // pack the inputs
+ for (int i = 0; i < inputs.size(); ++i) {
+ auto val = inputs.get(i);
+ if (val.isTuple()) {
+ auto tuples = val.toTuple()->elements();
+ for (auto& ival : tuples) {
+ TORCH_CHECK(ival.isTensor());
+ auto tensor = ival.toTensor();
+ PTMCoreMLFeatureSpecs spec{
+ .name = inputs_[inputSpecIndex].name(),
+ .tensor = tensor,
+ };
+ inputSpecs.emplace_back(spec);
+ ++inputSpecIndex;
+ }
+ } else {
+ TORCH_CHECK(val.isTensor());
+ auto tensor = val.toTensor();
+ PTMCoreMLFeatureSpecs spec{
+ .name = inputs_[inputSpecIndex].name(),
+ .tensor = tensor,
+ };
+ inputSpecs.emplace_back(spec);
+ ++inputSpecIndex;
+ }
+ }
+ // pack the outputs
+ c10::List<torch::Tensor> outputs;
+ id<MLFeatureProvider> results = [executor_ forwardWithInputs:inputSpecs];
+ for (auto& spec : outputs_) {
+ MLFeatureValue* val = [results featureValueForName:spec.name()];
+ TORCH_CHECK(val.multiArrayValue);
+ // Currently, only Float type is supported
+ TORCH_CHECK(val.multiArrayValue.dataType == MLMultiArrayDataTypeFloat32);
+ auto tensor = at::empty(spec.sizes(), scalarType(spec.dtype()));
+ int64_t count = val.multiArrayValue.count;
+ memcpy(
+ tensor.data_ptr<float>(),
+ (float*)val.multiArrayValue.dataPointer,
+ count * sizeof(float));
+ outputs.push_back(tensor);
+ }
+ return outputs;
+ }
+
+ private:
+ PTMCoreMLExecutor* executor_ = nullptr;
+ std::vector<TensorSpec> inputs_;
+ std::vector<TensorSpec> outputs_;
+ CoreMLConfig config_;
+};
+
+class API_AVAILABLE(ios(11.0), macos(10.13)) CoreMLBackend
+ : public torch::jit::PyTorchBackendInterface {
+ public:
+ c10::impl::GenericDict compile(
+ c10::IValue processed,
+ c10::impl::GenericDict method_compile_spec) override {
+ auto modelDict = processed.toGenericDict();
+ NSString* specs = [[NSString alloc]
+ initWithCString:modelDict.at("extra").toStringRef().c_str()
+ encoding:NSUTF8StringEncoding];
+ NSDictionary* dict = parse(specs);
+ NSArray<NSArray*>* inputs = dict[@"inputs"];
+ NSArray<NSArray*>* outputs = dict[@"outputs"];
+ std::vector<TensorSpec> inputSpecs, outputSpecs;
+ for (NSArray* input in inputs) {
+ inputSpecs.emplace_back(TensorSpec(input));
+ }
+ for (NSArray* output in outputs) {
+ outputSpecs.emplace_back(TensorSpec(output));
+ }
+ auto config = CoreMLConfig(dict[@"config"]);
+ const std::string& model = modelDict.at("model").toStringRef();
+ const std::string& sha256 = modelDict.at("hash").toStringRef();
+ PTMCoreMLExecutor* executor = [PTMCoreMLExecutor new];
+ bool result = [executor compileMLModel:model identifier:sha256];
+ TORCH_CHECK(result, "Compiling MLModel failed!");
+ auto executorWrapper = c10::make_intrusive<CoreMLExecutorWrapper>(
+ executor, inputSpecs, outputSpecs, config);
+ auto handle = IValue::make_capsule(executorWrapper);
+ c10::Dict<IValue, IValue> ret(StringType::get(), c10::AnyType::get());
+ ret.insert("forward", handle);
+ return c10::impl::toGenericDict(ret);
+ }
+
+ c10::impl::GenericList execute(
+ c10::IValue handle,
+ c10::impl::GenericList inputs) override {
+ auto executor = c10::static_intrusive_pointer_cast<CoreMLExecutorWrapper>(
+ handle.toCapsule());
+ auto outputs = executor->execute(inputs);
+ return c10::impl::toList(outputs);
+ }
+ bool is_available() override {
+#if !defined(__APPLE__)
+ return false;
+#else
+ if (@available(iOS 14, macOS 10.13, *)) {
+ return true;
+ } else {
+ return false;
+ }
+#endif
+ }
+};
+
+API_AVAILABLE(ios(11.0), macos(10.13))
+static auto cls = torch::jit::backend<CoreMLBackend>("coreml");
+
+} // namespace
+}
+}
+}
--- /dev/null
+#include <torch/csrc/jit/backends/coreml/objc/PTMCoreMLExecutor.h>
+#include <torch/script.h>
+
+#import <CoreML/CoreML.h>
+
+#include <sys/utsname.h>
+#include <fstream>
+#include <iostream>
+
+@implementation PTMCoreMLFeatureProvider {
+ NSUInteger _coremlVersion;
+ std::vector<PTMCoreMLFeatureSpecs> _specs;
+}
+
+@synthesize featureNames = _featureNames;
+
+- (instancetype)initWithFeatureSpecs:
+ (const std::vector<PTMCoreMLFeatureSpecs>&)specs
+ CoreMLVersion:(NSUInteger)ver {
+ self = [super init];
+ if (self) {
+ _coremlVersion = ver;
+ _specs = specs;
+ NSMutableArray* names = [NSMutableArray new];
+ for (auto& spec : _specs) {
+ [names addObject:spec.name];
+ }
+ _featureNames = [[NSSet alloc] initWithArray:names];
+ }
+ return self;
+}
+
+- (nullable MLFeatureValue*)featureValueForName:(NSString*)featureName {
+ for (auto& spec : _specs) {
+ if ([spec.name isEqualToString:featureName]) {
+ NSMutableArray* shape = [NSMutableArray new];
+ for (auto& dim : spec.tensor.sizes().vec()) {
+ [shape addObject:@(dim)];
+ }
+ NSMutableArray* strides = [NSMutableArray new];
+ for (auto& step : spec.tensor.strides().vec()) {
+ [strides addObject:@(step)];
+ }
+ NSError* error = nil;
+ TORCH_CHECK(spec.tensor.dtype() == c10::kFloat);
+ MLMultiArray* mlArray = [[MLMultiArray alloc]
+ initWithDataPointer:spec.tensor.data_ptr<float>()
+ shape:shape
+ dataType:MLMultiArrayDataTypeFloat32
+ strides:strides
+ deallocator:(^(void* bytes){
+ })error:&error];
+ return [MLFeatureValue featureValueWithMultiArray:mlArray];
+ }
+ }
+ return nil;
+}
+
+@end
+
+@implementation PTMCoreMLExecutor {
+ MLModel* _mlModel;
+}
+
+- (BOOL)compileMLModel:(const std::string&)modelSpecs
+ identifier:(const std::string&)identifier
+ API_AVAILABLE(ios(11.0), macos(10.13)) {
+ _modelPath = [self _save:modelSpecs
+ identifier:[NSString stringWithCString:identifier.c_str()
+ encoding:NSUTF8StringEncoding]];
+ NSError* error;
+ NSURL* compiledModelPath = nil;
+ if (@available(iOS 11.0, macOS 10.13, *)) {
+ compiledModelPath =
+ [MLModel compileModelAtURL:[NSURL URLWithString:_modelPath]
+ error:&error];
+ } else {
+ TORCH_CHECK(false, "CoreML is not available on your deivce");
+ }
+ if (error || !compiledModelPath) {
+ // remove cached models if compalition failed.
+ [self cleanup];
+ TORCH_CHECK(
+ false,
+ "Error compiling model",
+ [error localizedDescription].UTF8String);
+ return NO;
+ }
+ if (@available(iOS 12.0, macOS 10.14, *)) {
+ MLModelConfiguration* config = [MLModelConfiguration alloc];
+ MLComputeUnits backend = MLComputeUnitsCPUOnly;
+ if ([self.backend isEqualToString:@"cpuandgpu"]) {
+ backend = MLComputeUnitsCPUAndGPU;
+ } else if ([self.backend isEqualToString:@"all"]) {
+ backend = MLComputeUnitsAll;
+ }
+ config.computeUnits = backend;
+ config.allowLowPrecisionAccumulationOnGPU = self.allowLowPrecision;
+ _mlModel = [MLModel modelWithContentsOfURL:compiledModelPath
+ configuration:config
+ error:&error];
+ } else {
+ _mlModel = [MLModel modelWithContentsOfURL:compiledModelPath error:&error];
+ }
+ if (error || !_mlModel) {
+ TORCH_CHECK(
+ false, "Error loading MLModel", error.localizedDescription.UTF8String);
+ }
+
+ _compiledModelPath = compiledModelPath.path;
+ return YES;
+}
+
+- (id<MLFeatureProvider>)forwardWithInputs:
+ (const std::vector<PTMCoreMLFeatureSpecs>&)inputs {
+ NSError* error;
+ PTMCoreMLFeatureProvider* inputFeature = [[PTMCoreMLFeatureProvider alloc]
+ initWithFeatureSpecs:inputs
+ CoreMLVersion:self.coreMLVersion];
+ if (inputFeature == nil) {
+ NSLog(@"inputFeature is not initialized.");
+ return nil;
+ }
+ if (@available(iOS 11.0, macOS 10.13, *)) {
+ MLPredictionOptions* options = [[MLPredictionOptions alloc] init];
+ id<MLFeatureProvider> outputFeature =
+ [_mlModel predictionFromFeatures:inputFeature
+ options:options
+ error:&error];
+ if (error || !outputFeature) {
+ TORCH_CHECK(
+ false,
+ "Error running the prediction",
+ error.localizedDescription.UTF8String);
+ }
+
+ return outputFeature;
+ } else {
+ TORCH_CHECK("Core ML is available on iOS 11.0 and above");
+ return nil;
+ }
+}
+
+- (BOOL)cleanup {
+ NSFileManager* fileManager = [NSFileManager defaultManager];
+ NSError* error;
+ if (![fileManager fileExistsAtPath:_modelPath]) {
+ [fileManager removeItemAtPath:_modelPath error:&error];
+ }
+ if (![fileManager fileExistsAtPath:_compiledModelPath]) {
+ [fileManager removeItemAtPath:_compiledModelPath error:&error];
+ }
+ return !error;
+}
+
+- (NSString*)_save:(const std::string&)spec identifier:(NSString*)identifier {
+ NSURL* temporaryDirectoryURL = [NSURL fileURLWithPath:NSTemporaryDirectory()
+ isDirectory:YES];
+ NSString* modelPath = [NSString
+ stringWithFormat:@"%@/%@", temporaryDirectoryURL.path, identifier];
+ if (![[NSFileManager defaultManager] fileExistsAtPath:modelPath]) {
+ // Note that the serialized protobuf binary contains bytes, not text;
+ // see
+ // https://developers.google.com/protocol-buffers/docs/pythontutorial#parsing-and-serialization
+ NSData* data = [NSData dataWithBytes:spec.c_str() length:spec.length()];
+ BOOL ret = [data writeToFile:modelPath atomically:YES];
+ TORCH_CHECK(ret, "Save MLModel failed!", modelPath.UTF8String);
+ }
+ return modelPath;
+}
+
+@end
projects / platform / upstream / pytorch.git / commitdiff