提交 9fbd9426 编写于 作者: Q qijun

init

上级 7901f06a
/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserve.
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 "paddle/framework/op_registry.h"
#include "paddle/platform/place.h"
namespace paddle {
namespace operators {
class GetPlacesOp : public framework::OperatorBase {
public:
GetPlacesOp(const std::string &type, const framework::VariableNameMap &inputs,
const framework::VariableNameMap &outputs,
const framework::AttributeMap &attrs)
: OperatorBase(type, inputs, outputs, attrs) {}
void Run(const framework::Scope &scope,
const platform::DeviceContext &dev_ctx) const override {
auto use_gpu = Attr<bool>("use_gpu");
auto trainer_count = Attr<int>("trainer_count");
auto out_var_name = Output("Out");
auto *out_var = scope.FindVar(out_var_name);
PADDLE_ENFORCE(out_var != nullptr, "Output variable %s cannot be found",
out_var_name);
auto &places = *(out_var->GetMutable<std::vector<platform::Place>>());
places.reserve(trainer_count);
if (use_gpu) {
for (int i = 0; i < trainer_count; i++) {
places.emplace_back(platform::GPUPlace(i));
}
} else {
for (int i = 0; i < trainer_count; i++) {
places.emplace_back(platform::CPUPlace());
}
}
}
};
class GetPlacesOpProtoMaker : public framework::OpProtoAndCheckerMaker {
public:
GetPlacesOpProtoMaker(framework::OpProto *proto,
framework::OpAttrChecker *op_checker)
: OpProtoAndCheckerMaker(proto, op_checker) {
AddOutput("Out", "vector of Place");
AddAttr<int>("trainer_count", "(int)trainer count").SetDefault(1);
AddAttr<bool>("use_gpu", "(bool)use gpu").SetDefault(false);
AddComment(R"DOC(
GetPlaces Operator.
)DOC");
}
};
} // namespace operators
} // namespace paddle
namespace ops = paddle::operators;
REGISTER_OPERATOR(get_places, ops::GetPlacesOp, ops::GetPlacesOpProtoMaker);
...@@ -424,7 +424,7 @@ class Operator(object): ...@@ -424,7 +424,7 @@ class Operator(object):
self.desc.check_attrs() self.desc.check_attrs()
no_kernel_op_set = { no_kernel_op_set = {
'feed', 'fetch', 'save', 'load', 'recurrent', 'feed', 'fetch', 'save', 'load', 'recurrent',
'rnn_memory_helper_grad', 'conditional_block', 'while' 'rnn_memory_helper_grad', 'conditional_block', 'while', 'get_places'
} }
if type not in no_kernel_op_set: if type not in no_kernel_op_set:
self.desc.infer_var_type(self.block.desc) self.desc.infer_var_type(self.block.desc)
......
...@@ -8,6 +8,8 @@ import tensor ...@@ -8,6 +8,8 @@ import tensor
from tensor import * from tensor import *
import control_flow import control_flow
from control_flow import * from control_flow import *
import utils
from utils import *
__all__ = [] __all__ = []
__all__ += nn.__all__ __all__ += nn.__all__
...@@ -15,3 +17,4 @@ __all__ += io.__all__ ...@@ -15,3 +17,4 @@ __all__ += io.__all__
__all__ += tensor.__all__ __all__ += tensor.__all__
__all__ += control_flow.__all__ __all__ += control_flow.__all__
__all__ += ops.__all__ __all__ += ops.__all__
__all__ += utils.__all__
"""
All util layers.
"""
from ..layer_helper import LayerHelper
from ..framework import Variable
__all__ = ['get_places']
def get_places(use_gpu, trainer_count):
helper = LayerHelper('get_places', **locals())
out_places = helper.create_tmp_variable(dtype=helper.input_dtype())
helper.append_op(
type='get_places',
outputs={"Out": [out_places]},
attrs={
"use_gpu": use_gpu,
'trainer_count': trainer_count,
})
return out_places
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册