1 /** 2 * Copyright 2020-2021 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 #ifndef DATASET_ENGINE_DATASETOPS_EPOCH_CTRL_OP_H_ 17 #define DATASET_ENGINE_DATASETOPS_EPOCH_CTRL_OP_H_ 18 19 #include <memory> 20 #include <string> 21 #include <vector> 22 #include "minddata/dataset/engine/datasetops/repeat_op.h" 23 #include "minddata/dataset/engine/datasetops/pipeline_op.h" 24 25 namespace mindspore { 26 namespace dataset { 27 class EpochCtrlOp : public RepeatOp { 28 public: 29 // Constructor 30 explicit EpochCtrlOp(int32_t num_epoch); 31 32 // Destructor 33 ~EpochCtrlOp(); 34 35 // A print method typically used for debugging 36 // @param out - The output stream to write output to 37 // @param show_all - A bool to control if you want to show all info or just a summary 38 void Print(std::ostream &out, bool show_all) const override; Name()39 std::string Name() const override { return kEpochCtrlOp; } 40 41 // This function returns the row that is at the top of our output connector. The caller is 42 // typically our parent node, when the parent is asking us to provide the next row of data. 43 // Since EpochCtrlOp is derived from RepeatOp which is an inlined op, getting a row from us 44 // will simply bounce you to get a row from our child. 45 // Epoch Control Op does not eat the EOE, it will pass the EOE to the next op. 46 Status GetNextRow(TensorRow *row, int32_t worker_id, bool retry_if_eoe) override; 47 48 // Base-class override for handling cases when an eoe is received. 49 // @param worker_id - The worker id 50 Status EoeReceived(int32_t worker_id) override; 51 52 int64_t GetTreeRepeatCount() override; 53 }; 54 } // namespace dataset 55 } // namespace mindspore 56 57 #endif // DATASET_ENGINE_DATASETOPS_EPOCH_CTRL_OP_H_ 58