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,
