mindspore-ci-bot
5 years ago
committed by
Gitee
commit
5c42198430
@ -0,0 +1,63 @@
|
||||
/**
|
||||
* Copyright 2020 Huawei Technologies Co., Ltd
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
#ifndef MINDSPORE_LITE_INCLUDE_TRAIN_SESSION_H_
|
||||
#define MINDSPORE_LITE_INCLUDE_TRAIN_SESSION_H_
|
||||
#include <vector>
|
||||
#include <string>
|
||||
#include <unordered_map>
|
||||
// #include "include/lite_session.h"
|
||||
#include "src/lite_session.h"
|
||||
|
||||
namespace mindspore {
|
||||
namespace lite {
|
||||
class Model;
|
||||
}
|
||||
namespace lite::tensor {
|
||||
class Tensor;
|
||||
}
|
||||
namespace session {
|
||||
|
||||
class TrainSession : public lite::LiteSession {
|
||||
public:
|
||||
TrainSession();
|
||||
~TrainSession() = default;
|
||||
|
||||
int RunGraph(const session::KernelCallBack &before = nullptr,
|
||||
const session::KernelCallBack &after = nullptr) override;
|
||||
|
||||
int CompileGraph(lite::Model *model) override;
|
||||
virtual void ReplaceOps();
|
||||
virtual void* ExportToBuf(void* buf, size_t* len) const;
|
||||
|
||||
std::unordered_map<std::string, std::vector<mindspore::tensor::MSTensor *>> GetOutputs() const;
|
||||
std::vector<tensor::MSTensor *> GetOutputsByName(const std::string &node_name) const;
|
||||
|
||||
virtual void train();
|
||||
bool is_train() { return train_mode_ == true; }
|
||||
virtual void eval();
|
||||
bool is_eval() { return train_mode_ == false; }
|
||||
|
||||
protected:
|
||||
bool train_mode_ = false;
|
||||
lite::Model* model_ = nullptr;
|
||||
std::unordered_map<std::string, std::vector<mindspore::tensor::MSTensor *>> ext_output_map_;
|
||||
|
||||
|
||||
// private:
|
||||
};
|
||||
} // namespace session
|
||||
} // namespace mindspore
|
||||
#endif // MINDSPORE_LITE_INCLUDE_TRAIN_SESSION_H_
|
||||
@ -0,0 +1,64 @@
|
||||
/**
|
||||
* Copyright 2019-2020 Huawei Technologies Co., Ltd
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
#include "src/ops/apply_momentum.h"
|
||||
namespace mindspore {
|
||||
namespace lite {
|
||||
|
||||
|
||||
#ifdef PRIMITIVE_WRITEABLE
|
||||
|
||||
#else
|
||||
int ApplyMomentum::UnPackToFlatBuilder(const schema::Primitive *primitive, flatbuffers::FlatBufferBuilder *fbb) {
|
||||
MS_ASSERT(nullptr != primitive);
|
||||
MS_ASSERT(nullptr != fbb);
|
||||
auto attr = primitive->value_as_ApplyMomentum();
|
||||
if (attr == nullptr) {
|
||||
MS_LOG(ERROR) << "value_as_ApplyMomentum return nullptr";
|
||||
return RET_ERROR;
|
||||
}
|
||||
auto val_offset = schema::CreateApplyMomentum(*fbb);
|
||||
auto prim_offset = schema::CreatePrimitive(*fbb, schema::PrimitiveType_ActivationGrad, val_offset.o);
|
||||
fbb->Finish(prim_offset);
|
||||
return RET_OK;
|
||||
}
|
||||
#endif
|
||||
|
||||
int ApplyMomentum::InferShape(std::vector<tensor::Tensor *> inputs, std::vector<tensor::Tensor *> outputs) {
|
||||
if (5 != inputs.size()) {
|
||||
MS_LOG(ERROR) << "ApplyMomentum should have at 5 input tensors";
|
||||
return RET_ERROR;
|
||||
}
|
||||
// if (outputs.empty()) {
|
||||
// MS_LOG(ERROR) << "ApplyMomentumCPUKernel error input output size!";
|
||||
// return RET_ERROR;
|
||||
// }
|
||||
|
||||
if (inputs[0]->ElementsNum() != inputs[1]->ElementsNum() || inputs[0]->ElementsNum() != inputs[3]->ElementsNum() ||
|
||||
inputs[2]->ElementsNum() != 1 || inputs[4]->ElementsNum() != 1) {
|
||||
MS_LOG(ERROR) << "error input data size!";
|
||||
return RET_ERROR;
|
||||
}
|
||||
if (!outputs.empty()) {
|
||||
auto *out = outputs.front();
|
||||
MS_ASSERT(out != nullptr);
|
||||
out->set_data_type(inputs[0]->data_type());
|
||||
out->SetFormat(inputs[0]->GetFormat());
|
||||
}
|
||||
|
||||
return RET_OK;
|
||||
}
|
||||
} // namespace lite
|
||||
} // namespace mindspore
|
||||
@ -0,0 +1,44 @@
|
||||
/**
|
||||
* Copyright 2019-2020 Huawei Technologies Co., Ltd
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
#ifndef MINDSPORE_LITE_SRC_OPS_APPLY_MOMENTUM_H_
|
||||
#define MINDSPORE_LITE_SRC_OPS_APPLY_MOMENTUM_H_
|
||||
|
||||
#include <vector>
|
||||
#include <set>
|
||||
#include <cmath>
|
||||
#include "ir/dtype/type_id.h"
|
||||
#include "src/ops/primitive_c.h"
|
||||
|
||||
namespace mindspore {
|
||||
namespace lite {
|
||||
class ApplyMomentum : public PrimitiveC {
|
||||
public:
|
||||
#ifdef PRIMITIVE_WRITEABLE
|
||||
MS_DECLARE_PARENT(ApplyMomentum, PrimitiveC);
|
||||
ApplyMomentum() = default;
|
||||
explicit ApplyMomentum(schema::PrimitiveT *primitive) : PrimitiveC(primitive) {}
|
||||
#else
|
||||
ApplyMomentum() = default;
|
||||
|
||||
int UnPackToFlatBuilder(const schema::Primitive *primitive, flatbuffers::FlatBufferBuilder *fbb) override;
|
||||
#endif
|
||||
int InferShape(std::vector<lite::tensor::Tensor *> inputs_, std::vector<lite::tensor::Tensor *> outputs_) override;
|
||||
};
|
||||
} // namespace lite
|
||||
} // namespace mindspore
|
||||
|
||||
#endif // MINDSPORE_LITE_SRC_OPS_APPLY_MOMENTUM_H_
|
||||
Some files were not shown because too many files have changed in this diff Show More
Loading…
Reference in new issue