def _forward_grouped_convolution_xp(self, x, gy, xp): G = self.groups N, iC = x.shape[:2] oC = gy.shape[1] o_size = gy.shape[2:] o_size_prod = utils.size_of_shape(o_size) k_size = self.ksize dims = len(o_size) iCg = iC // G oCg = oC // G # Do not check iCg and oCg because this class is rarely used alone # (N, iC, k_size..., o_size...) x = conv_nd.im2col_nd(x, k_size, self.stride, self.pad, cover_all=self.cover_all, dilate=self.dilate) x = xp.rollaxis(x, 0, dims + 2) # (iC, k_size..., N, o_size...) mul_len = iCg * utils.size_of_shape(k_size) x = x.reshape(G, mul_len, N * o_size_prod) x = x.transpose(0, 2, 1) # (G, N*o_size, iCg*k_size) gy = xp.rollaxis(gy, 1) # (oC, N, o_size...) gy = gy.reshape(G, oCg, N * o_size_prod) # (G, oCg, iCg*k_size) = (G, oCg, N*o_size) @ (G, N*o_size, iCg*k_size) gW = convolution_2d._matmul(gy, x).astype(self.W_dtype, copy=False) gW = gW.reshape(oC, iCg, *k_size) return gW,
def _forward_grouped_convolution_xp(self, x, W, b, xp): # G: group count # N: batch size # xC: input channels # yC: output channels G = self.groups N, xC = x.shape[:2] x_size = x.shape[2:] yCg = W.shape[1] yC = yCg * G xCg = xC // G k_size = W.shape[2:] dims = len(k_size) if xC % G != 0: raise TypeError('The number of groups must be ' 'a divisor of that of input channels') x = xp.rollaxis(x, 1) # (xC, N, x_size...) x = x.reshape(G, xCg, N * convolution_nd._prod(x_size)) W = W.reshape(G, xCg, yCg * convolution_nd._prod(k_size)) W = W.transpose(0, 2, 1) # (G, yCg*k_size, xCg) # (G, yCg*k_size, N*x_size) = (G, yCg*k_size, xCg) @ (G, xCg, N*x_size) col = convolution_2d._matmul(W, x).astype(x.dtype, copy=False) col = col.reshape((yC,) + k_size + (N,) + x_size) col = xp.rollaxis(col, dims + 1) # (N, yC, k_size..., x_size...) y = conv_nd.col2im_nd(col, self.stride, self.pad, self.outs, dilate=self.dilate) if b is not None: y += b.reshape(1, yC, *((1,) * dims)) return y,
def _forward_grouped_convolution(self, x, W, b): # G: group count # N: batch size # kH, kW: kernel height, kernel width # xC, xH, xW: x channels, x height, x width # yC, yH, yW: y channels, y height, y width G = self.groups N, xC, xH, xW = x.shape xCg = xC // G _, yCg, kH, kW = W.shape # _ == xC yC = yCg * G x = x.transpose(1, 0, 2, 3) # (xC, N, xH, xW) x = x.reshape(G, xCg, N * xH * xW) W = W.reshape(G, xCg, yCg * kH * kW) W = W.transpose(0, 2, 1) # (G, yCg*kH*kW, xCg) # (G, yCg*kH*kW, N*xH*xW) = (G, yCg*kH*kW, xCg) @ (G, xCg, N*xH*xW) col = convolution_2d._matmul(W, x).astype(x.dtype, copy=False) col = col.reshape(yC, kH, kW, N, xH, xW) col = col.transpose(3, 0, 1, 2, 4, 5) # (N, yC, kH, kW, xH, xW) y = conv.col2im(col, self.sy, self.sx, self.ph, self.pw, self.outh, self.outw, dy=self.dy, dx=self.dx) if b is not None: y += b.reshape(1, b.size, 1, 1) return y,
def _forward_grouped_convolution_xp(self, x, gy, xp): G = self.groups N, iC = x.shape[:2] oC = gy.shape[1] o_size = gy.shape[2:] o_size_prod = utils.size_of_shape(o_size) k_size = self.ksize dims = len(o_size) iCg = iC // G oCg = oC // G # Do not check iCg and oCg because this class is rarely used alone # (N, iC, k_size..., o_size...) x = conv_nd.im2col_nd(x, k_size, self.stride, self.pad, cover_all=self.cover_all, dilate=self.dilate) x = xp.rollaxis(x, 0, dims + 2) # (iC, k_size..., N, o_size...) mul_len = iCg * utils.size_of_shape(k_size) x = x.reshape(G, mul_len, N * o_size_prod) x = x.transpose(0, 2, 1) # (G, N*o_size, iCg*k_size) gy = xp.rollaxis(gy, 1) # (oC, N, o_size...) gy = gy.reshape(G, oCg, N * o_size_prod) # (G, oCg, iCg*k_size) = (G, oCg, N*o_size) @ (G, N*o_size, iCg*k_size) gW = convolution_2d._matmul(gy, x).astype(self.W_dtype, copy=False) gW = gW.reshape(oC, iCg, *k_size) return gW,
def _forward_grouped_convolution_xp(self, x, W, b, xp): # G: group count # N: batch size # iC: input channels # oC: output channels G = self.groups N, iC = x.shape[:2] oC = W.shape[0] k_size = W.shape[2:] iCg = iC // G oCg = oC // G dims = len(k_size) if iC % G != 0: raise TypeError('The number of groups must be ' 'a divisor of that of input channels') if oC % G != 0: raise TypeError('The number of groups must be ' 'a divisor of that of output channels') xp = backend.get_array_module(x) # (N, iC, k_size..., o_size...) x = conv_nd.im2col_nd(x, k_size, self.stride, self.pad, cover_all=self.cover_all, dilate=self.dilate) o_size = x.shape[-dims:] x = xp.rollaxis(x, 0, dims + 2) # (iC, k_size..., N, o_size...) mul_len = iCg * utils.size_of_shape(k_size) x = x.reshape(G, mul_len, N * utils.size_of_shape(o_size)) W = W.reshape(G, oCg, mul_len) # (G, oCg, N*o_size) = (G, oCg, iCg*k_size) @ (G, iCg*k_size, N*o_size) y = convolution_2d._matmul(W, x).astype(x.dtype, copy=False) y = y.reshape(oC, N, *o_size) y = xp.rollaxis(y, 1) # (N, oC, o_size...) if b is not None: y += b.reshape(1, b.size, *((1, ) * dims)) return y,
def _forward_grouped_convolution_xp(self, x, W, b, xp): # G: group count # N: batch size # iC: input channels # oC: output channels G = self.groups N, iC = x.shape[:2] oC = W.shape[0] k_size = W.shape[2:] iCg = iC // G oCg = oC // G dims = len(k_size) if iC % G != 0: raise TypeError('The number of groups must be ' 'a divisor of that of input channels') if oC % G != 0: raise TypeError('The number of groups must be ' 'a divisor of that of output channels') xp = backend.get_array_module(x) # (N, iC, k_size..., o_size...) x = conv_nd.im2col_nd(x, k_size, self.stride, self.pad, cover_all=self.cover_all, dilate=self.dilate) o_size = x.shape[-dims:] x = xp.rollaxis(x, 0, dims + 2) # (iC, k_size..., N, o_size...) mul_len = iCg * utils.size_of_shape(k_size) x = x.reshape(G, mul_len, N * utils.size_of_shape(o_size)) W = W.reshape(G, oCg, mul_len) # (G, oCg, N*o_size) = (G, oCg, iCg*k_size) @ (G, iCg*k_size, N*o_size) y = convolution_2d._matmul(W, x).astype(x.dtype, copy=False) y = y.reshape(oC, N, *o_size) y = xp.rollaxis(y, 1) # (N, oC, o_size...) if b is not None: y += b.reshape(1, b.size, *((1,) * dims)) return y,
def _forward_grouped_convolution_xp(self, x, W, b, xp): # G: group count # N: batch size # xC: input channels # yC: output channels G = self.groups N, xC = x.shape[:2] x_size = x.shape[2:] yCg = W.shape[1] yC = yCg * G xCg = xC // G k_size = W.shape[2:] dims = len(k_size) if xC % G != 0: raise TypeError('The number of groups must be ' 'a divisor of that of input channels') x = xp.rollaxis(x, 1) # (xC, N, x_size...) x = x.reshape(G, xCg, N * utils.size_of_shape(x_size)) W = W.reshape(G, xCg, yCg * utils.size_of_shape(k_size)) W = W.transpose(0, 2, 1) # (G, yCg*k_size, xCg) # (G, yCg*k_size, N*x_size) = (G, yCg*k_size, xCg) @ (G, xCg, N*x_size) col = convolution_2d._matmul(W, x).astype(x.dtype, copy=False) col = col.reshape((yC, ) + k_size + (N, ) + x_size) col = xp.rollaxis(col, dims + 1) # (N, yC, k_size..., x_size...) y = conv_nd.col2im_nd(col, self.stride, self.pad, self.outs, dilate=self.dilate) if b is not None: y += b.reshape(1, yC, *((1, ) * dims)) return y,