wehub-resource-sync
78 lines
2.4 KiB
C++
78 lines
2.4 KiB
C++
//
|
|
// TorchZeros.cpp
|
|
// MNNConverter
|
|
//
|
|
// Created by MNN on 2022/04/28.
|
|
// Copyright © 2018, Alibaba Group Holding Limited
|
|
//
|
|
|
|
#include "MNN_generated.h"
|
|
#include "TorchExtraManager.hpp"
|
|
#include "logkit.h"
|
|
|
|
namespace MNN {
|
|
namespace Express {
|
|
|
|
class TorchZerosTransform : public TorchExtraManager::Transform {
|
|
public:
|
|
virtual EXPRP onExecute(EXPRP expr) const override {
|
|
auto inputs = expr->inputs();
|
|
auto op = expr->get();
|
|
auto opName = op->name()->str();
|
|
MNN_ASSERT(inputs.size() == 1);
|
|
auto zeros = _Fill(inputs[0], _Const(0.));
|
|
zeros->setName(opName);
|
|
return zeros->expr().first;
|
|
}
|
|
};
|
|
|
|
class TorchOnesTransform : public TorchExtraManager::Transform {
|
|
public:
|
|
virtual EXPRP onExecute(EXPRP expr) const override {
|
|
auto inputs = expr->inputs();
|
|
auto op = expr->get();
|
|
auto opName = op->name()->str();
|
|
MNN_ASSERT(inputs.size() == 1);
|
|
auto zeros = _Fill(inputs[0], _Const(1.));
|
|
zeros->setName(opName);
|
|
return zeros->expr().first;
|
|
}
|
|
};
|
|
|
|
class TorchOnesLikeTransform : public TorchExtraManager::Transform {
|
|
public:
|
|
virtual EXPRP onExecute(EXPRP expr) const override {
|
|
auto inputs = expr->inputs();
|
|
auto op = expr->get();
|
|
auto opName = op->name()->str();
|
|
MNN_ASSERT(inputs.size() == 1);
|
|
auto zeros = _Fill(_Shape(inputs[0]), _Const(1.));
|
|
zeros->setName(opName);
|
|
return zeros->expr().first;
|
|
}
|
|
};
|
|
|
|
class TorchFullLikeTransform : public TorchExtraManager::Transform {
|
|
public:
|
|
virtual EXPRP onExecute(EXPRP expr) const override {
|
|
auto inputs = expr->inputs();
|
|
auto op = expr->get();
|
|
auto opName = op->name()->str();
|
|
MNN_ASSERT(inputs.size() == 2);
|
|
auto full = _Fill(_Shape(inputs[0]), inputs[1]);
|
|
full->setName(opName);
|
|
return full->expr().first;
|
|
}
|
|
};
|
|
|
|
static auto gRegister = []() {
|
|
TorchExtraManager::get()->insert("zeros", std::shared_ptr<TorchExtraManager::Transform>(new TorchZerosTransform));
|
|
TorchExtraManager::get()->insert("ones", std::shared_ptr<TorchExtraManager::Transform>(new TorchOnesTransform));
|
|
TorchExtraManager::get()->insert("ones_like", std::shared_ptr<TorchExtraManager::Transform>(new TorchOnesLikeTransform));
|
|
TorchExtraManager::get()->insert("full_like", std::shared_ptr<TorchExtraManager::Transform>(new TorchFullLikeTransform));
|
|
return true;
|
|
}();
|
|
|
|
} // namespace Express
|
|
} // namespace MNN
|