1 /** 2 * Copyright 2020 Huawei Technologies Co., Ltd 3 * 4 * Licensed under the Apache License, Version 2.0 (the "License"); 5 * you may not use this file except in compliance with the License. 6 * You may obtain a copy of the License at 7 * 8 * http://www.apache.org/licenses/LICENSE-2.0 9 * 10 * Unless required by applicable law or agreed to in writing, software 11 * distributed under the License is distributed on an "AS IS" BASIS, 12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 * See the License for the specific language governing permissions and 14 * limitations under the License. 15 */ 16 17 #ifndef MINDSPORE_CCSRC_MINDDATA_DATASET_ENGINE_IR_DATASETOPS_SOURCE_TRANSFER_NODE_H_ 18 #define MINDSPORE_CCSRC_MINDDATA_DATASET_ENGINE_IR_DATASETOPS_SOURCE_TRANSFER_NODE_H_ 19 20 #include <memory> 21 #include <string> 22 #include <vector> 23 24 #include "minddata/dataset/engine/ir/datasetops/dataset_node.h" 25 26 namespace mindspore { 27 namespace dataset { 28 29 class TransferNode : public DatasetNode { 30 public: 31 /// \brief Constructor 32 TransferNode(std::shared_ptr<DatasetNode> child, std::string queue_name, std::string device_type, int32_t device_id, 33 bool send_epoch_end, int32_t total_batch, bool create_data_info_queue); 34 35 /// \brief Destructor 36 ~TransferNode() = default; 37 38 /// \brief Node name getter 39 /// \return Name of the current node Name()40 std::string Name() const override { return kTransferNode; } 41 42 /// \brief Print the description 43 /// \param out - The output stream to write output to 44 void Print(std::ostream &out) const override; 45 46 /// \brief Copy the node to a new object 47 /// \return A shared pointer to the new copy 48 std::shared_ptr<DatasetNode> Copy() override; 49 50 /// \brief a base class override function to create the required runtime dataset op objects for this class 51 /// \param node_ops - A vector containing shared pointer to the Dataset Ops that this object will create 52 /// \return Status Status::OK() if build successfully 53 Status Build(std::vector<std::shared_ptr<DatasetOp>> *const node_ops) override; 54 55 /// \brief Parameters validation 56 /// \return Status Status::OK() if all the parameters are valid 57 Status ValidateParams() override; 58 59 static Status get_distribution(std::shared_ptr<DatasetNode> ds, int32_t *device_id); 60 61 /// \brief Base-class override for accepting IRNodePass visitor 62 /// \param[in] p The node to visit 63 /// \param[out] modified Indicator if the node was modified 64 /// \return Status of the node visit 65 Status Accept(IRNodePass *const p, bool *const modified) override; 66 67 /// \brief Base-class override for accepting IRNodePass visitor 68 /// \param[in] p The node to visit 69 /// \param[out] modified Indicator if the node was modified 70 /// \return Status of the node visit 71 Status AcceptAfter(IRNodePass *const p, bool *const modified) override; 72 73 /// \brief Getter functions QueueName()74 const std::string &QueueName() const { return queue_name_; } DeviceId()75 int32_t DeviceId() const { return device_id_; } DeviceType()76 const std::string &DeviceType() const { return device_type_; } PrefetchSize()77 int32_t PrefetchSize() const { return prefetch_size_; } SendEpochEnd()78 bool SendEpochEnd() const { return send_epoch_end_; } TotalBatch()79 int32_t TotalBatch() const { return total_batch_; } CreateDataInfoQueue()80 bool CreateDataInfoQueue() const { return create_data_info_queue_; } 81 82 /// \brief Get the arguments of node 83 /// \param[out] out_json JSON string of all attributes 84 /// \return Status of the function 85 Status to_json(nlohmann::json *out_json) override; 86 87 /// \brief Function for read dataset operation from json 88 /// \param[in] json_obj The JSON object to be deserialized 89 /// \param[in] ds dataset node constructed 90 /// \param[out] result Deserialized dataset after the operation 91 /// \return Status The status code returned 92 static Status from_json(nlohmann::json json_obj, std::shared_ptr<DatasetNode> ds, 93 std::shared_ptr<DatasetNode> *result); 94 95 private: 96 std::string queue_name_; 97 int32_t device_id_; 98 std::string device_type_; 99 int32_t prefetch_size_; 100 bool send_epoch_end_; 101 int32_t total_batch_; 102 bool create_data_info_queue_; 103 }; 104 105 } // namespace dataset 106 } // namespace mindspore 107 #endif // MINDSPORE_CCSRC_MINDDATA_DATASET_ENGINE_IR_DATASETOPS_SOURCE_TRANSFER_NODE_H_ 108