save + load model using checkpoints not working

[brettkoonce]

save + load model using checkpoints not working

[brettkoonce]

working here, have trained a resnet50 model using checkpoints, thanks again for the notes!

[BradLarson]

This is due to a shortcoming in the current checkpoint reader's iteration implementation. When I first put this together, there wasn't a way to use reflection to determine if a property was writeable, so it was possible to have the reflection part of the iteration find a path that shouldn't be part of the checkpoint. As a hacky way around this, I let you specify properties that you can ignore for reading and writing via ignoredTensorPaths in the Checkpointable protocol.

In the case of ResNet (and anything else that uses BatchNorm), inside of BatchNorm are the properties runningMean and runningVariance that are not to be used inside a checkpoint. You need to manually ignore these by adjusting the conformance to

extension ResNet: Checkpointable {
  public var ignoredTensorPaths: Set<String> {
    return ["BatchNorm<Float>.runningMean", "BatchNorm<Float>.runningVariance"]
  }
}

The good news is that the new reflection enhancements upstream will allow for immutable properties to be determined, and this workaround can go away. However, because that's not supported in all of our toolchains, I held off on reworking this to use that interface.

Additionally, there's also a current flaw in the checkpoint reading where you'll need to change your code above to read

try model.readCheckpoint(from: temporaryDirectory.appendingPathComponent("ResNet"), name: "ResNet")

for it to find the written checkpoint. I need to align the directory naming between reading and writing checkpoints to do away with this inconsistency. Sorry about the implementation flaws here.

[brettkoonce]

@BradLarson thanks for the notes, that works here!