393 lines
15 KiB
Plaintext
393 lines
15 KiB
Plaintext
// Copyright 2022 Google Inc. All rights reserved.
|
|
//
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
// you may not use this file except in compliance with the License.
|
|
// You may obtain a copy of the License at:
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
// See the License for the specific language governing permissions and
|
|
// limitations under the License.
|
|
|
|
#import "TFLSignatureRunner+Internal.h"
|
|
|
|
#include <vector>
|
|
|
|
#import "TFLCommonUtil.h"
|
|
#import "TFLErrorUtil.h"
|
|
#import "TFLInterpreter+Internal.h"
|
|
#import "TFLQuantizationParameters+Internal.h"
|
|
#import "TFLTensor+Internal.h"
|
|
|
|
#ifdef COCOAPODS
|
|
#import <TensorFlowLiteC/TensorFlowLiteC.h>
|
|
#else
|
|
#include "tensorflow/lite/core/c/c_api_experimental.h"
|
|
#endif // COCOAPODS
|
|
|
|
NS_ASSUME_NONNULL_BEGIN
|
|
|
|
/** Domain for errors in the signature runner. */
|
|
NSErrorDomain const TFLSignatureRunnerErrorDomain = @"org.tensorflow.lite.SignatureRunner";
|
|
|
|
@interface TFLSignatureRunner ()
|
|
|
|
/**
|
|
* The backing interpreter. It's a strong reference to ensure that the interpreter is never released
|
|
* before this signature runner is released.
|
|
*
|
|
* @warning Never let the interpreter hold a strong reference to the signature runner to avoid
|
|
* retain cycles.
|
|
*/
|
|
@property(nonatomic, readonly) TFLInterpreter *interpreter;
|
|
|
|
/** TfLiteSignatureRunner backed by C API. */
|
|
@property(nonatomic, readonly) TfLiteSignatureRunner *signatureRunner;
|
|
|
|
@end
|
|
|
|
@implementation TFLSignatureRunner {
|
|
// Whether we need to allocate tensors memory.
|
|
BOOL _isTensorsAllocationNeeded;
|
|
}
|
|
|
|
@synthesize inputs = _inputs;
|
|
@synthesize outputs = _outputs;
|
|
@synthesize signatureKey = _signatureKey;
|
|
|
|
#pragma mark - Initializer
|
|
|
|
- (nullable instancetype)initWithInterpreter:(TFLInterpreter *)interpreter
|
|
signatureKey:(NSString *)signatureKey
|
|
error:(NSError **)error {
|
|
self = [super init];
|
|
if (self != nil) {
|
|
_signatureKey = [signatureKey copy];
|
|
const char *signatureKeyCString = _signatureKey.UTF8String;
|
|
TfLiteSignatureRunner *signatureRunner =
|
|
TfLiteInterpreterGetSignatureRunner(interpreter.interpreter, signatureKeyCString);
|
|
if (signatureRunner == nullptr) {
|
|
NSString *errorDescription =
|
|
[NSString stringWithFormat:
|
|
@"Failed to create a signature runner. Signature with key (%@) not found.",
|
|
signatureKey];
|
|
[TFLErrorUtil setError:error
|
|
withDomain:TFLSignatureRunnerErrorDomain
|
|
code:TFLSignatureRunnerErrorCodeFailedToCreateSignatureRunner
|
|
description:errorDescription];
|
|
return nil;
|
|
}
|
|
_signatureRunner = signatureRunner;
|
|
_interpreter = interpreter;
|
|
_isTensorsAllocationNeeded = YES;
|
|
[self allocateTensorsWithError:error];
|
|
}
|
|
return self;
|
|
}
|
|
|
|
- (void)dealloc {
|
|
TfLiteSignatureRunnerDelete(_signatureRunner);
|
|
}
|
|
|
|
#pragma mark - Public
|
|
|
|
- (NSArray<NSString *> *)inputs {
|
|
if (_inputs) return _inputs;
|
|
NSUInteger inputCount = TfLiteSignatureRunnerGetInputCount(self.signatureRunner);
|
|
NSMutableArray<NSString *> *mutableInputsArray =
|
|
[[NSMutableArray alloc] initWithCapacity:inputCount];
|
|
for (NSUInteger i = 0; i < inputCount; i++) {
|
|
const char *inputNameCString =
|
|
TfLiteSignatureRunnerGetInputName(self.signatureRunner, (int32_t)i);
|
|
NSString *inputName = @"";
|
|
if (inputNameCString != nullptr) {
|
|
inputName = [NSString stringWithUTF8String:inputNameCString] ?: @"";
|
|
};
|
|
[mutableInputsArray addObject:inputName];
|
|
}
|
|
_inputs = [mutableInputsArray copy];
|
|
return _inputs;
|
|
}
|
|
|
|
- (NSArray<NSString *> *)outputs {
|
|
if (_outputs) return _outputs;
|
|
NSUInteger outputCount = TfLiteSignatureRunnerGetOutputCount(self.signatureRunner);
|
|
NSMutableArray<NSString *> *mutableOutputsArray =
|
|
[[NSMutableArray alloc] initWithCapacity:outputCount];
|
|
for (NSUInteger i = 0; i < outputCount; i++) {
|
|
const char *outputNameCString =
|
|
TfLiteSignatureRunnerGetOutputName(self.signatureRunner, (int32_t)i);
|
|
NSString *outputName = @"";
|
|
if (outputNameCString != nullptr) {
|
|
outputName = [NSString stringWithUTF8String:outputNameCString] ?: @"";
|
|
}
|
|
[mutableOutputsArray addObject:outputName];
|
|
}
|
|
_outputs = [mutableOutputsArray copy];
|
|
return _outputs;
|
|
}
|
|
|
|
- (nullable TFLTensor *)inputTensorWithName:(NSString *)name error:(NSError **)error {
|
|
return [self tensorOfType:TFLTensorTypeInput nameInSignature:name error:error];
|
|
}
|
|
|
|
- (nullable TFLTensor *)outputTensorWithName:(NSString *)name error:(NSError **)error {
|
|
return [self tensorOfType:TFLTensorTypeOutput nameInSignature:name error:error];
|
|
}
|
|
|
|
- (BOOL)resizeInputTensorWithName:(NSString *)name
|
|
toShape:(NSArray<NSNumber *> *)shape
|
|
error:(NSError **)error {
|
|
if (shape.count == 0) {
|
|
[TFLErrorUtil setError:error
|
|
withDomain:TFLSignatureRunnerErrorDomain
|
|
code:TFLSignatureRunnerErrorCodeInvalidShape
|
|
description:@"Invalid shape. Must not be empty."];
|
|
return NO;
|
|
}
|
|
|
|
std::vector<int> cDimensions(shape.count);
|
|
for (int dimIndex = 0; dimIndex < shape.count; ++dimIndex) {
|
|
int dimension = shape[dimIndex].intValue;
|
|
if (dimension <= 0) {
|
|
[TFLErrorUtil setError:error
|
|
withDomain:TFLSignatureRunnerErrorDomain
|
|
code:TFLSignatureRunnerErrorCodeInvalidShape
|
|
description:@"Invalid shape. Dimensions must be positive integers."];
|
|
return NO;
|
|
}
|
|
cDimensions[dimIndex] = dimension;
|
|
}
|
|
|
|
if (TfLiteSignatureRunnerResizeInputTensor(self.signatureRunner, name.UTF8String,
|
|
cDimensions.data(),
|
|
(int32_t)shape.count) != kTfLiteOk) {
|
|
NSString *errorDescription =
|
|
[NSString stringWithFormat:@"Failed to resize input tensor with input name (%@).", name];
|
|
[TFLErrorUtil setError:error
|
|
withDomain:TFLSignatureRunnerErrorDomain
|
|
code:TFLSignatureRunnerErrorCodeFailedToResizeInputTensor
|
|
description:errorDescription];
|
|
return NO;
|
|
}
|
|
|
|
// Need to reallocate tensor memory.
|
|
_isTensorsAllocationNeeded = YES;
|
|
|
|
return YES;
|
|
}
|
|
|
|
- (BOOL)allocateTensorsWithError:(NSError **)error {
|
|
if (!_isTensorsAllocationNeeded) return YES;
|
|
if (TfLiteSignatureRunnerAllocateTensors(self.signatureRunner) != kTfLiteOk) {
|
|
[TFLErrorUtil setError:error
|
|
withDomain:TFLSignatureRunnerErrorDomain
|
|
code:TFLSignatureRunnerErrorCodeFailedToAllocateTensors
|
|
description:@"Failed to allocate memory for tensors."];
|
|
return NO;
|
|
}
|
|
_isTensorsAllocationNeeded = NO;
|
|
return YES;
|
|
}
|
|
|
|
- (BOOL)invokeWithInputs:(NSDictionary<NSString *, NSData *> *)inputs Error:(NSError **)error {
|
|
if (![self allocateTensorsWithError:error]) return NO;
|
|
|
|
// Fill in input data.
|
|
for (NSString *inputName in inputs.allKeys) {
|
|
TFLTensor *inputTensor = [self inputTensorWithName:inputName error:error];
|
|
if (!inputTensor) return NO;
|
|
if (![inputTensor copyData:inputs[inputName] error:error]) return NO;
|
|
}
|
|
|
|
if (TfLiteSignatureRunnerInvoke(self.signatureRunner) != kTfLiteOk) {
|
|
[TFLErrorUtil setError:error
|
|
withDomain:TFLSignatureRunnerErrorDomain
|
|
code:TFLSignatureRunnerErrorCodeFailedToInvoke
|
|
description:@"Failed to invoke the signature runner."];
|
|
return NO;
|
|
}
|
|
|
|
return YES;
|
|
}
|
|
|
|
#pragma mark - TFLTensorDataAccessor
|
|
|
|
- (BOOL)copyData:(NSData *)data toInputTensor:(TFLTensor *)inputTensor error:(NSError **)error {
|
|
if (inputTensor.type == TFLTensorTypeOutput) {
|
|
[TFLErrorUtil setError:error
|
|
withDomain:TFLSignatureRunnerErrorDomain
|
|
code:TFLSignatureRunnerErrorCodeCopyDataToOutputTensorNotAllowed
|
|
description:@"Cannot copy data into an output tensor."];
|
|
return NO;
|
|
}
|
|
const TfLiteTensor *cTensor = [self cTensorOfType:TFLTensorTypeInput
|
|
nameInSignature:inputTensor.nameInSignature
|
|
error:error];
|
|
if (cTensor == nullptr) {
|
|
return NO;
|
|
}
|
|
|
|
NSUInteger byteSize = (NSUInteger)TfLiteTensorByteSize(cTensor);
|
|
if (data.length != byteSize) {
|
|
NSString *errorDescription = [NSString
|
|
stringWithFormat:
|
|
@"Input tensor with input name (%@) expects data size (%lu), but got (%lu).",
|
|
inputTensor.nameInSignature, (unsigned long)byteSize, (unsigned long)data.length];
|
|
[TFLErrorUtil setError:error
|
|
withDomain:TFLSignatureRunnerErrorDomain
|
|
code:TFLSignatureRunnerErrorCodeInvalidInputByteSize
|
|
description:errorDescription];
|
|
return NO;
|
|
}
|
|
|
|
if (TfLiteTensorCopyFromBuffer((TfLiteTensor *)cTensor, data.bytes, data.length) != kTfLiteOk) {
|
|
NSString *errorDescription =
|
|
[NSString stringWithFormat:@"Failed to copy data into input tensor with input name (%@).",
|
|
inputTensor.nameInSignature];
|
|
[TFLErrorUtil setError:error
|
|
withDomain:TFLSignatureRunnerErrorDomain
|
|
code:TFLSignatureRunnerErrorCodeFailedToCopyDataToInputTensor
|
|
description:errorDescription];
|
|
return NO;
|
|
}
|
|
|
|
return YES;
|
|
}
|
|
|
|
- (nullable NSData *)dataFromTensor:(TFLTensor *)tensor error:(NSError **)error {
|
|
const TfLiteTensor *cTensor = [self cTensorOfType:tensor.type
|
|
nameInSignature:tensor.nameInSignature
|
|
error:error];
|
|
if (cTensor == nullptr) {
|
|
return nil;
|
|
}
|
|
|
|
void *bytes = TfLiteTensorData(cTensor);
|
|
NSUInteger byteSize = (NSUInteger)TfLiteTensorByteSize(cTensor);
|
|
if (bytes == nullptr || byteSize == 0) {
|
|
NSString *tensorType = [TFLTensor stringForTensorType:tensor.type];
|
|
NSString *errorDescription =
|
|
[NSString stringWithFormat:@"Failed to get data from %@ tensor with %@ name (%@).",
|
|
tensorType, tensorType, tensor.nameInSignature];
|
|
[TFLErrorUtil setError:error
|
|
withDomain:TFLSignatureRunnerErrorDomain
|
|
code:TFLSignatureRunnerErrorCodeFailedToGetDataFromTensor
|
|
description:errorDescription];
|
|
return nil;
|
|
}
|
|
|
|
return [NSData dataWithBytes:bytes length:byteSize];
|
|
}
|
|
|
|
- (nullable NSArray<NSNumber *> *)shapeOfTensor:(TFLTensor *)tensor error:(NSError **)error {
|
|
const TfLiteTensor *cTensor = [self cTensorOfType:tensor.type
|
|
nameInSignature:tensor.nameInSignature
|
|
error:error];
|
|
if (cTensor == nullptr) {
|
|
return nil;
|
|
}
|
|
|
|
NSString *tensorType = [TFLTensor stringForTensorType:tensor.type];
|
|
int32_t rank = TfLiteTensorNumDims(cTensor);
|
|
if (rank <= 0) {
|
|
NSString *errorDescription =
|
|
[NSString stringWithFormat:@"%@ tensor with %@ name (%@) has invalid rank (%d).",
|
|
tensorType, tensorType, tensor.nameInSignature, rank];
|
|
[TFLErrorUtil setError:error
|
|
withDomain:TFLSignatureRunnerErrorDomain
|
|
code:TFLSignatureRunnerErrorCodeInvalidTensor
|
|
description:errorDescription];
|
|
return nil;
|
|
}
|
|
|
|
NSMutableArray<NSNumber *> *shape = [NSMutableArray arrayWithCapacity:rank];
|
|
for (int32_t dimIndex = 0; dimIndex < rank; dimIndex++) {
|
|
int32_t dimension = TfLiteTensorDim(cTensor, dimIndex);
|
|
if (dimension <= 0) {
|
|
NSString *errorDescription = [NSString
|
|
stringWithFormat:@"%@ tensor with %@ name (%@) has invalid %d-th dimension (%d).",
|
|
tensorType, tensorType, tensor.nameInSignature, dimIndex, dimension];
|
|
[TFLErrorUtil setError:error
|
|
withDomain:TFLSignatureRunnerErrorDomain
|
|
code:TFLSignatureRunnerErrorCodeInvalidTensor
|
|
description:errorDescription];
|
|
return nil;
|
|
}
|
|
shape[dimIndex] = @((NSUInteger)dimension);
|
|
}
|
|
|
|
return shape;
|
|
}
|
|
|
|
#pragma mark - Private
|
|
|
|
- (nullable TFLTensor *)tensorOfType:(TFLTensorType)type
|
|
nameInSignature:(NSString *)nameInSignature
|
|
error:(NSError **)error {
|
|
const TfLiteTensor *tensor = [self cTensorOfType:type
|
|
nameInSignature:nameInSignature
|
|
error:error];
|
|
|
|
if (tensor == nullptr) {
|
|
return nil;
|
|
}
|
|
|
|
NSString *tensorName = TFLTensorNameFromCTensor(tensor);
|
|
if (!tensorName) {
|
|
NSString *tensorType = [TFLTensor stringForTensorType:type];
|
|
NSString *errorDescription =
|
|
[NSString stringWithFormat:@"Failed to get name of %@ tensor with %@ name (%@).",
|
|
tensorType, tensorType, nameInSignature];
|
|
[TFLErrorUtil setError:error
|
|
withDomain:TFLSignatureRunnerErrorDomain
|
|
code:TFLSignatureRunnerErrorCodeInvalidTensor
|
|
description:errorDescription];
|
|
return nil;
|
|
}
|
|
TFLTensorDataType dataType = TFLTensorDataTypeFromCTensor(tensor);
|
|
TFLQuantizationParameters *quantizationParams = TFLQuantizationParamsFromCTensor(tensor);
|
|
|
|
return [[TFLTensor alloc] initWithSignatureRunner:self
|
|
type:type
|
|
nameInSignature:nameInSignature
|
|
name:tensorName
|
|
dataType:dataType
|
|
quantizationParameters:quantizationParams];
|
|
}
|
|
|
|
- (const TfLiteTensor *)cTensorOfType:(TFLTensorType)type
|
|
nameInSignature:(NSString *)nameInSignature
|
|
error:(NSError **)error {
|
|
const TfLiteTensor *tensor = nullptr;
|
|
const char *nameCString = nameInSignature.UTF8String;
|
|
switch (type) {
|
|
case TFLTensorTypeInput:
|
|
tensor = TfLiteSignatureRunnerGetInputTensor(self.signatureRunner, nameCString);
|
|
break;
|
|
case TFLTensorTypeOutput:
|
|
tensor = TfLiteSignatureRunnerGetOutputTensor(self.signatureRunner, nameCString);
|
|
break;
|
|
}
|
|
|
|
if (tensor == nullptr) {
|
|
NSString *tensorType = [TFLTensor stringForTensorType:type];
|
|
NSString *errorDescription =
|
|
[NSString stringWithFormat:@"Failed to get %@ tensor with %@ name (%@).", tensorType,
|
|
tensorType, nameInSignature];
|
|
[TFLErrorUtil setError:error
|
|
withDomain:TFLSignatureRunnerErrorDomain
|
|
code:TFLSignatureRunnerErrorCodeFailedToGetTensor
|
|
description:errorDescription];
|
|
}
|
|
|
|
return tensor;
|
|
}
|
|
|
|
@end
|
|
|
|
NS_ASSUME_NONNULL_END
|