tfnetwork.h 3.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103
  1. ///////////////////////////////////////////////////////////////////////
  2. // File: tfnetwork.h
  3. // Description: Encapsulation of an entire tensorflow graph as a
  4. // Tesseract Network.
  5. // Author: Ray Smith
  6. // Created: Fri Feb 26 09:35:29 PST 2016
  7. //
  8. // (C) Copyright 2016, Google Inc.
  9. // Licensed under the Apache License, Version 2.0 (the "License");
  10. // you may not use this file except in compliance with the License.
  11. // You may obtain a copy of the License at
  12. // http://www.apache.org/licenses/LICENSE-2.0
  13. // Unless required by applicable law or agreed to in writing, software
  14. // distributed under the License is distributed on an "AS IS" BASIS,
  15. // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  16. // See the License for the specific language governing permissions and
  17. // limitations under the License.
  18. ///////////////////////////////////////////////////////////////////////
  19. #ifndef TESSERACT_LSTM_TFNETWORK_H_
  20. #define TESSERACT_LSTM_TFNETWORK_H_
  21. #ifdef INCLUDE_TENSORFLOW
  22. #include <memory>
  23. #include <string>
  24. #include "network.h"
  25. #include "static_shape.h"
  26. #include "tfnetwork.pb.h"
  27. #include "tensorflow/core/framework/graph.pb.h"
  28. #include "tensorflow/core/public/session.h"
  29. namespace tesseract {
  30. class TFNetwork : public Network {
  31. public:
  32. explicit TFNetwork(const STRING& name);
  33. virtual ~TFNetwork() = default;
  34. // Returns the required shape input to the network.
  35. StaticShape InputShape() const override { return input_shape_; }
  36. // Returns the shape output from the network given an input shape (which may
  37. // be partially unknown ie zero).
  38. StaticShape OutputShape(const StaticShape& input_shape) const override {
  39. return output_shape_;
  40. }
  41. STRING spec() const override { return spec_.c_str(); }
  42. // Deserializes *this from a serialized TFNetwork proto. Returns 0 if failed,
  43. // otherwise the global step of the serialized graph.
  44. int InitFromProtoStr(const std::string& proto_str);
  45. // The number of classes in this network should be equal to those in the
  46. // recoder_ in LSTMRecognizer.
  47. int num_classes() const { return output_shape_.depth(); }
  48. // Writes to the given file. Returns false in case of error.
  49. // Should be overridden by subclasses, but called by their Serialize.
  50. bool Serialize(TFile* fp) const override;
  51. // Reads from the given file. Returns false in case of error.
  52. // Should be overridden by subclasses, but NOT called by their DeSerialize.
  53. bool DeSerialize(TFile* fp) override;
  54. // Runs forward propagation of activations on the input line.
  55. // See Network for a detailed discussion of the arguments.
  56. void Forward(bool debug, const NetworkIO& input,
  57. const TransposedArray* input_transpose,
  58. NetworkScratch* scratch, NetworkIO* output) override;
  59. private:
  60. // Runs backward propagation of errors on the deltas line.
  61. // See Network for a detailed discussion of the arguments.
  62. bool Backward(bool debug, const NetworkIO& fwd_deltas,
  63. NetworkScratch* scratch,
  64. NetworkIO* back_deltas) override {
  65. tprintf("Must override Network::DebugWeights for type %d\n", type_);
  66. }
  67. void DebugWeights() override {
  68. tprintf("Must override Network::DebugWeights for type %d\n", type_);
  69. }
  70. int InitFromProto();
  71. // The original network definition for reference.
  72. std::string spec_;
  73. // Input tensor parameters.
  74. StaticShape input_shape_;
  75. // Output tensor parameters.
  76. StaticShape output_shape_;
  77. // The tensor flow graph is contained in here.
  78. std::unique_ptr<tensorflow::Session> session_;
  79. // The serialized graph is also contained in here.
  80. TFNetworkModel model_proto_;
  81. };
  82. } // namespace tesseract.
  83. #endif // ifdef INCLUDE_TENSORFLOW
  84. #endif // TESSERACT_TENSORFLOW_TFNETWORK_H_