Skip to content
Permalink
Branch: master
Find file Copy path
Find file Copy path
Fetching contributors…
Cannot retrieve contributors at this time
196 lines (164 sloc) 6.58 KB
/*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you 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.
*/
/*!
* Copyright (c) 2015 by Contributors
* \file sequence_mask.cc
* \brief
* \author Sebastian Bodenstein
*/
#include "./sequence_mask-inl.h"
namespace mxnet {
namespace op {
// (seqlen, batch, rest) case
template <int req>
struct SequenceMask0CPUKernel {
template <typename DType, typename IType>
MSHADOW_XINLINE static void Map(int batch, DType *in, const IType *idx,
index_t max_s_len, index_t batch_size,
index_t restsize, DType value) {
const index_t seqpos = static_cast<int>(idx[batch]);
#pragma unroll
for (index_t s = seqpos; s < max_s_len; ++s) {
index_t incr = (s * batch_size * restsize) + (batch * restsize);
#pragma unroll
for (index_t r = 0; r < restsize; ++r)
KERNEL_ASSIGN(in[incr + r], req, value);
}
}
};
// (batch, seqlen, rest) case
template <int req>
struct SequenceMask1CPUKernel {
template <typename DType, typename IType>
MSHADOW_XINLINE static void Map(int batch, DType *in, const IType *idx,
index_t max_s_len, index_t batch_size,
index_t restsize, DType value) {
const index_t seqpos = static_cast<int>(idx[batch]);
#pragma unroll
for (index_t s = seqpos; s < max_s_len; ++s) {
index_t incr = (batch * max_s_len * restsize) + (s * restsize);
#pragma unroll
for (index_t r = 0; r < restsize; ++r)
KERNEL_ASSIGN(in[incr + r], req, value);
}
}
};
template<typename DType, typename IType>
void SequenceMaskExec(
const mshadow::Tensor<cpu, 3, DType> &data,
const mshadow::Tensor<cpu, 1, IType> &indices,
const OpReqType req, mshadow::Stream<cpu> *const s,
int axis, DType val) {
using namespace mshadow;
using namespace mshadow::expr;
using namespace mxnet_op;
index_t batch = indices.size(0);
index_t max_seq_len = data.size(axis);
index_t restsize = data.size(2);
MXNET_ASSIGN_REQ_SWITCH(req, req_type, {
if (axis == 1) {
Kernel<SequenceMask1CPUKernel<req_type>, cpu>::Launch(
s, batch, data.dptr_, indices.dptr_, max_seq_len, batch, restsize,
val);
} else {
Kernel<SequenceMask0CPUKernel<req_type>, cpu>::Launch(
s, batch, data.dptr_, indices.dptr_, max_seq_len, batch, restsize,
val);
}
});
}
template <>
Operator *CreateOp<cpu>(SequenceMaskParam param, int dtype, int itype) {
Operator *op = nullptr;
MSHADOW_TYPE_SWITCH(dtype, DType, {
MSHADOW_TYPE_SWITCH(itype, IType, {
op = new SequenceMaskOp<cpu, DType, IType>(param);
});
});
return op;
}
// DO_BIND_DISPATCH comes from operator_common.h
Operator *SequenceMaskProp::CreateOperatorEx(Context ctx,
mxnet::ShapeVector *in_shape,
std::vector<int> *in_type) const {
if (in_type->size() >= 2 && (*in_type)[1] != -1) {
DO_BIND_DISPATCH(CreateOp, param_, (*in_type)[0], (*in_type)[1]);
}
// sequence_length not passed in, so fall back to using input array dtype for second argument
DO_BIND_DISPATCH(CreateOp, param_, (*in_type)[0], (*in_type)[0]);
}
DMLC_REGISTER_PARAMETER(SequenceMaskParam);
MXNET_REGISTER_OP_PROPERTY(SequenceMask, SequenceMaskProp)
.describe(R"code(Sets all elements outside the sequence to a constant value.
This function takes an n-dimensional input array of the form
[max_sequence_length, batch_size, other_feature_dims] and returns an array of the same shape.
Parameter `sequence_length` is used to handle variable-length sequences. `sequence_length`
should be an input array of positive ints of dimension [batch_size].
To use this parameter, set `use_sequence_length` to `True`,
otherwise each example in the batch is assumed to have the max sequence length and
this operator works as the `identity` operator.
Example::
x = [[[ 1., 2., 3.],
[ 4., 5., 6.]],
[[ 7., 8., 9.],
[ 10., 11., 12.]],
[[ 13., 14., 15.],
[ 16., 17., 18.]]]
// Batch 1
B1 = [[ 1., 2., 3.],
[ 7., 8., 9.],
[ 13., 14., 15.]]
// Batch 2
B2 = [[ 4., 5., 6.],
[ 10., 11., 12.],
[ 16., 17., 18.]]
// works as identity operator when sequence_length parameter is not used
SequenceMask(x) = [[[ 1., 2., 3.],
[ 4., 5., 6.]],
[[ 7., 8., 9.],
[ 10., 11., 12.]],
[[ 13., 14., 15.],
[ 16., 17., 18.]]]
// sequence_length [1,1] means 1 of each batch will be kept
// and other rows are masked with default mask value = 0
SequenceMask(x, sequence_length=[1,1], use_sequence_length=True) =
[[[ 1., 2., 3.],
[ 4., 5., 6.]],
[[ 0., 0., 0.],
[ 0., 0., 0.]],
[[ 0., 0., 0.],
[ 0., 0., 0.]]]
// sequence_length [2,3] means 2 of batch B1 and 3 of batch B2 will be kept
// and other rows are masked with value = 1
SequenceMask(x, sequence_length=[2,3], use_sequence_length=True, value=1) =
[[[ 1., 2., 3.],
[ 4., 5., 6.]],
[[ 7., 8., 9.],
[ 10., 11., 12.]],
[[ 1., 1., 1.],
[ 16., 17., 18.]]]
)code" ADD_FILELINE)
.add_argument("data", "NDArray-or-Symbol",
"n-dimensional input array of the form [max_sequence_length,"
" batch_size, other_feature_dims] where n>2")
.add_argument("sequence_length", "NDArray-or-Symbol",
"vector of sequence lengths of the form [batch_size]")
.add_arguments(SequenceMaskParam::__FIELDS__());
} // namespace op
} // namespace mxnet
You can’t perform that action at this time.