Added checkpointing

apple · Oct 3, 2019 · fb8c383 · fb8c383
1 parent f10ccb2
commit fb8c383
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 6 deletions.
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
@@ -104,13 +104,13 @@ - (instancetype) initWithParams:(NSString *)name
       size_t offset = index * _numberOfFeatureChannels * sizeof(float);
 
       style_property.gammaBuffer = [dev newBufferWithBytes:(char *) _gamma_weights.bytes + offset
-                                                    length:sizeof(float) * _numberOfFeatureChannels
-                                                   options:MTLResourceStorageModeManaged];
-
-      style_property.betaBuffer = [dev newBufferWithBytes:(char *) _beta_weights.bytes + offset
                                                    length:sizeof(float) * _numberOfFeatureChannels
                                                   options:MTLResourceStorageModeManaged];
 
+      style_property.betaBuffer = [dev newBufferWithBytes:(char *) _beta_weights.bytes + offset
+                                                  length:sizeof(float) * _numberOfFeatureChannels
+                                                 options:MTLResourceStorageModeManaged];
+
       style_property.gammaMomentumBuffer = [dev newBufferWithBytes:zeros_ptr
                                                             length:sizeof(float) * _numberOfFeatureChannels
                                                            options:MTLResourceStorageModeManaged];
@@ -196,8 +196,8 @@ - (void) loadGamma:(float *)gamma {
 
 // TODO: refactor for multiple indicies
 - (float *) gamma {
-  NSUInteger previousStyle = _styleIndex;
   _gammaPlaceHolder = [NSMutableData data];
+  NSUInteger previousStyle = _styleIndex;
   for (NSUInteger index = 0; index < _styles; index++) { 
     _styleIndex = index; 
     [self checkpointWithCommandQueue:_cq];
@@ -243,7 +243,7 @@ - (MPSCNNNormalizationGammaAndBetaState *)updateGammaAndBetaWithCommandBuffer:(i
 
   }
 
-    return [[_style_props objectAtIndex: _styleIndex] state];
+  return [[_style_props objectAtIndex: _styleIndex] state];
 }
 
 - (void)checkpointWithCommandQueue:(nonnull id<MTLCommandQueue>)commandQueue {

diff --git a/src/ml/neural_net/style_transfer/mps_style_transfer.h b/src/ml/neural_net/style_transfer/mps_style_transfer.h
@@ -35,6 +35,7 @@ API_AVAILABLE(macos(10.15))
 - (NSDictionary<NSString *, NSData *> *) predict:(NSDictionary<NSString *, NSData *> *)inputs;
 - (void) setLearningRate:(float)lr;
 - (NSDictionary<NSString *, NSData *> *) train:(NSDictionary<NSString *, NSData *> *)inputs;
+- (void) checkpoint;
 
 @end
 

diff --git a/src/ml/neural_net/style_transfer/mps_style_transfer.m b/src/ml/neural_net/style_transfer/mps_style_transfer.m
@@ -252,6 +252,7 @@ - (instancetype) initWithDev:(id<MTLDevice>) dev
 }
 
 - (NSDictionary<NSString *, NSData *> *) exportWeights {
+  [self checkpoint];
   return [_model exportWeights:@"transformer_"];
 }
 
@@ -437,6 +438,14 @@ - (void) setLearningRate:(float)lr {
   return [lossDict copy];
 }
 
+- (void) checkpoint {
+  _inferenceGraph = [MPSNNGraph graphWithDevice:_dev
+                                    resultImage:_model.forwardPass
+                            resultImageIsNeeded:YES];
+
+  _inferenceGraph.format = MPSImageFeatureChannelFormatFloat32;
+}
+
 @end
 
 #endif // #ifdef HAS_MACOS_10_15