diff --git a/src/ml/neural_net/mps_layer_instance_norm_data_loader.mm b/src/ml/neural_net/mps_layer_instance_norm_data_loader.mm index 60e49c4439..613596baf8 100644 --- a/src/ml/neural_net/mps_layer_instance_norm_data_loader.mm +++ b/src/ml/neural_net/mps_layer_instance_norm_data_loader.mm @@ -177,11 +177,11 @@ - (void) loadBeta:(float *)beta { - (float *) beta { NSUInteger previousStyle = _styleIndex; - _betaPlaceHolder = [NSMutableData data]; + _betaPlaceHolder.length = 0; for (NSUInteger index = 0; index < _styles; index++) { _styleIndex = index; [self checkpointWithCommandQueue:_cq]; - float* betaWeights = (float *) [[[_style_props objectAtIndex: _styleIndex] betaBuffer] contents]; + float* betaWeights = (float *) [_style_props[_styleIndex].betaBuffer contents]; [_betaPlaceHolder appendBytes:betaWeights length:sizeof(float)*_numberOfFeatureChannels]; } _styleIndex = previousStyle; @@ -196,12 +196,12 @@ - (void) loadGamma:(float *)gamma { // TODO: refactor for multiple indicies - (float *) gamma { - _gammaPlaceHolder = [NSMutableData data]; NSUInteger previousStyle = _styleIndex; + _gammaPlaceHolder.length = 0; for (NSUInteger index = 0; index < _styles; index++) { _styleIndex = index; [self checkpointWithCommandQueue:_cq]; - float* gammaWeights = (float *) [[[_style_props objectAtIndex: _styleIndex] gammaBuffer] contents]; + float* gammaWeights = (float *) [_style_props[_styleIndex].gammaBuffer contents]; [_gammaPlaceHolder appendBytes:gammaWeights length:sizeof(float)*_numberOfFeatureChannels]; } _styleIndex = previousStyle; diff --git a/src/ml/neural_net/style_transfer/mps_style_transfer.m b/src/ml/neural_net/style_transfer/mps_style_transfer.m index 0517e62e4f..e74dff351e 100644 --- a/src/ml/neural_net/style_transfer/mps_style_transfer.m +++ b/src/ml/neural_net/style_transfer/mps_style_transfer.m @@ -438,6 +438,10 @@ - (void) setLearningRate:(float)lr { return [lossDict copy]; } +/** +* HACK: this somehow checkpoints the model for weight exports and updating the +* data loaders. Following up internally for a proper fix to this issue. +**/ - (void) checkpoint { _inferenceGraph = [MPSNNGraph graphWithDevice:_dev resultImage:_model.forwardPass diff --git a/src/python/turicreate/toolkits/style_transfer/style_transfer.py b/src/python/turicreate/toolkits/style_transfer/style_transfer.py index 843287c34c..61665da47b 100644 --- a/src/python/turicreate/toolkits/style_transfer/style_transfer.py +++ b/src/python/turicreate/toolkits/style_transfer/style_transfer.py @@ -409,6 +409,18 @@ def create(style_dataset, content_dataset, style_feature=None, mps_mxnet_key_map = _MpsStyleGraphAPI.mps_mxnet_weight_dict() + # LOGIC + # + # The two layers we are concerned about are very different. The instance + # norm layer weights should have a dimensionality two. Wheras the + # Convolutional Weights have a dimensionality of four. The Convolutional + # weights are in a different format in MxNet then they are in MPS. Since + # the arrays come back flattened in the MPS a reshape has to occur. But + # this reshape happens before the transpose therefore the shape itself + # has to be transposed. For the InstanceNorm Layer, however, the weights + # are passed back to MxNet in the correct format so just a simple + # reshape suffices. + # for key in mps_weights: if "transformer" in key: weight = transformer.collect_params()[mps_mxnet_key_map[key]].data()