From 72b829dcca2d1acaeea130e580ce780b1a7d550a Mon Sep 17 00:00:00 2001 From: Adrian Kuegel Date: Thu, 16 Aug 2018 01:32:25 -0700 Subject: Add a feature_group_size parameter to the Convolution HLO op. This is a first step towards supporting grouped convolutions, which are a generalization of depthwise convolution. PiperOrigin-RevId: 208950311 --- .../performance/xla/operation_semantics.md | 43 ++++++++++++++++------ 1 file changed, 31 insertions(+), 12 deletions(-) (limited to 'tensorflow/docs_src') diff --git a/tensorflow/docs_src/performance/xla/operation_semantics.md b/tensorflow/docs_src/performance/xla/operation_semantics.md index e24a7cda73..8c9d26fcbb 100644 --- a/tensorflow/docs_src/performance/xla/operation_semantics.md +++ b/tensorflow/docs_src/performance/xla/operation_semantics.md @@ -505,16 +505,17 @@ Computes a convolution of the kind used in neural networks. Here, a convolution can be thought of as a n-dimensional window moving across a n-dimensional base area and a computation is performed for each possible position of the window. -| Arguments | Type | Semantics | -| ---------------- | ----------------------- | ----------------------------- | -| `lhs` | `XlaOp` | rank n+2 array of inputs | -| `rhs` | `XlaOp` | rank n+2 array of kernel | -: : : weights : -| `window_strides` | `ArraySlice` | n-d array of kernel strides | -| `padding` | `ArraySlice>` : padding : -| `lhs_dilation` | `ArraySlice` | n-d lhs dilation factor array | -| `rhs_dilation` | `ArraySlice` | n-d rhs dilation factor array | +| Arguments | Type | Semantics | +| --------------------- | -------------------- | ----------------------------- | +| `lhs` | `XlaOp` | rank n+2 array of inputs | +| `rhs` | `XlaOp` | rank n+2 array of kernel | +: : : weights : +| `window_strides` | `ArraySlice` | n-d array of kernel strides | +| `padding` | `ArraySlice< | n-d array of (low, high) | +: : pair>` : padding : +| `lhs_dilation` | `ArraySlice` | n-d lhs dilation factor array | +| `rhs_dilation` | `ArraySlice` | n-d rhs dilation factor array | +| `feature_group_count` | int64 | the number of feature groups | Let n be the number of spatial dimensions. The `lhs` argument is a rank n+2 array describing the base area. This is called the input, even though of course @@ -532,8 +533,8 @@ The `rhs` argument is a rank n+2 array describing the convolutional filter/kernel/window. The dimensions are, in this order: * `output-z`: The `z` dimension of the output. -* `input-z`: The size of this dimension should equal the size of the `z` - dimension in lhs. +* `input-z`: The size of this dimension times `feature_group_count` should + equal the size of the `z` dimension in lhs. * `spatial_dims`: Describes the `n` spatial dimensions that define the n-d window that moves across the base area. @@ -566,6 +567,24 @@ Dilation of the rhs is also called atrous convolution. For more details, see `tf.nn.atrous_conv2d`. Dilation of the lhs is also called transposed convolution. For more details, see `tf.nn.conv2d_transpose`. +The `feature_group_count` argument (default value 1) can be used for grouped +convolutions. `feature_group_count` needs to be a divisor of both the input and +the output feature dimension. If `feature_group_count` is greater than 1, it +means that conceptually the input and output feature dimension and the `rhs` +output feature dimension are split evenly into `feature_group_count` many +groups, each group consisting of a consecutive subsequence of features. The +input feature dimension of `rhs` needs to be equal to the `lhs` input feature +dimension divided by `feature_group_count` (so it already has the size of a +group of input features). The i-th groups are used together to compute +`feature_group_count` many separate convolutions. The results of these +convolutions are concatenated together in the output feature dimension. + +For depthwise convolution the `feature_group_count` argument would be set to the +input feature dimension, and the filter would be reshaped from +`[filter_height, filter_width, in_channels, channel_multiplier]` to +`[filter_height, filter_width, 1, in_channels * channel_multiplier]`. For more +details, see `tf.nn.depthwise_conv2d`. + The output shape has these dimensions, in this order: * `batch`: Same size as `batch` on the input (`lhs`). -- cgit v1.2.3