com.intel.analytics.bigdl.nn.Mean.scala Maven / Gradle / Ivy
/*
* Copyright 2016 The BigDL Authors.
*
* 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.
*/
package com.intel.analytics.bigdl.nn
import com.intel.analytics.bigdl.tensor.TensorNumericMath.TensorNumeric
import scala.reflect.ClassTag
/**
* It is a simple layer which applies a mean operation over the given dimension.
* When nInputDims is provided, the input will be considered as batches.
* Then the mean operation will be applied in (dimension + 1).
*
* The input to this layer is expected to be a tensor, or a batch of tensors;
* when using mini-batch, a batch of sample tensors will be passed to the layer and
* the user need to specify the number of dimensions of each sample tensor in the
* batch using `nInputDims`.
*
* @param dimension the dimension to be applied mean operation
* @param nInputDims specify the number of dimensions that this module will receive
* If it is more than the dimension of input tensors, the first dimension
* would be considered as batch size
* @param squeeze default is true, which will squeeze the sum dimension; set it to false to keep
* the sum dimension
*/
@SerialVersionUID(2995626598003841724L)
class Mean[T: ClassTag](
val dimension: Int = 1,
val nInputDims: Int = -1,
val squeeze: Boolean = true)
(implicit ev: TensorNumeric[T])
extends Sum[T](dimension, nInputDims, true, squeeze) {
override def toString: String = s"nn.Mean"
}
object Mean {
def apply[@specialized(Float, Double) T: ClassTag](
dimension: Int = 1,
nInputDims: Int = -1,
squeeze: Boolean = true)(implicit ev: TensorNumeric[T]) : Mean[T] = {
new Mean[T](dimension, nInputDims, squeeze)
}
}