CODEWITHSUNDEEP
tensorflow
Tianjin Gu
Tianjin Gu

Reputation: 784

How to save and restore ResourceBase?

I'm add a self-defined ops and using ResourceMgr to a create child class of ResourceBase to store some status, and how to write the status into checkpoint file.

class FeatureTransformMap : public ResourceBase {
public:
    FeatureTransformMap(int32_t max_feature_count) : max_feature_count_(max_feature_count), feature_index_(0) {
        cout<<"Max feature count is:"<<max_feature_count_<<endl;
    }
    string DebugString() {return "FeatureTransformMap";}

    int32_t GetFeatureIndex(const string& feature) {
        {
            mutex_lock l(mu_);
            feature_index_ += 1;
        }
        return feature_index_;
    }



private:
    tensorflow::mutex mu_;
    uint32_t feature_index_ GUARDED_BY(mu_);
    const uint32_t max_feature_count_;
};

as in code above, how can I write feature_index_ to checkpoint files.

Upvotes: 0

Views: 240

Answers (1)

mrry
mrry

Reputation: 126184

There is no generic way to save a tensorflow::ResourceBase instance, but you can implement your own checkpointing support as follows:

  1. Define methods on your FeatureTransformMap class that serialize the state of the map to, and deserialize it from, one or more tensorflow::Tensor objects. For an example, see MutableHashTableOfScalars::ExportValues() and MutableHashTableOfScalars::ImportValues().

  2. Implement new TensorFlow OpKernel classes that invoke your serialization and deserialization methods. For an example, see LookupTableExportOp and LookupTableImportOp.

  3. In Python, implement a subclass of BaseSaverBuilder.SaveableObject for your resource, which includes invocations of your new ops. For an example, see MutableDenseHashTable._Saveable.

  4. In Python, when you create an instance of your resource, add it to the collection of tf.GraphKeys.SAVEABLE_OBJECTS. For an example, see how this is done for MutableDenseHashTable here.

Upvotes: 2

Related Questions