def space_to_patch(x, patchsz, paddings=None):
"""
[batch, imgsz0, imgsz1, depth]
=> [newbatch, patchsz0, patchsz1, depth]
=> [batch x blocks, imgsz0 / blocksz0, imgsz1 / blocksz1, depth]
:param tf.Tensor x:
:param patchsz:
:param paddings: [(top, bottom), (left, right)] if None then when need, auto paddings
:return:
"""
if paddings is not None:
if len(paddings) == 2:
paddings = [(0, 0), paddings[0], paddings[1], (0, 0)]
x = tf.pad(x, paddings)
inshape = x.dims
batch, imgsz, depth = inshape[0], inshape[1:3], inshape[3]
patchsz = [patchsz[0] or imgsz[0], patchsz[1] or imgsz[1]]
blocksz = (imgsz[0] // patchsz[0], imgsz[1] // patchsz[1])
if paddings is None:
# add padding if necessary
pd = [
imgsz[0] - blocksz[0] * patchsz[0],
imgsz[1] - blocksz[1] * patchsz[1]
]
if pd[0] or pd[1]:
pd = [patchsz[0] - pd[0], patchsz[1] - pd[1]]
paddings = [(0, 0), (pd[0] // 2, pd[0] - pd[0] // 2),
(pd[1] // 2, pd[1] - pd[1] // 2), (0, 0)]
x = tf.pad(x, paddings)
inshape = x.dims
batch, imgsz, depth = inshape[0], inshape[1:3], inshape[3]
blocksz = (imgsz[0] // patchsz[0], imgsz[1] // patchsz[1])
assert blocksz[0] * patchsz[0] == imgsz[0]
assert blocksz[1] * patchsz[1] == imgsz[1]
# batch *= (blocksz[0] * blocksz[1])
# [batch, imgsz0, imgsz1, d] => [batch, block0, patch0, block1, patch1, d]
data = x.reshape(batch or -1, blocksz[0], patchsz[0], blocksz[1],
patchsz[1], depth)
# dim to [batch, block0, block1, patch0, patch1, d]
data = data.transpose(0, 1, 3, 2, 4, 5)
# new batch size
if batch is not None:
batch *= (blocksz[0] * blocksz[1])
data = data.reshape(batch or -1, patchsz[0], patchsz[1], depth)
return data
def dwconv(x,
kernel,
multiplier=1,
stride=1,
pad=0,
padding='SAME',
initializer=tf.he_uniform,
bias=False,
**kwargs):
if pad:
pads = [(0, 0), (pad, pad), (pad, pad), (0, 0)]
x = tf.pad(x, pads, mode='CONSTANT')
kernel = _kernel_shape(2, kernel, x.dims[-1], multiplier)
stride = _stride_shape(2, stride)
W = tf.get_weight('W',
shape=kernel,
initializer=initializer(kernel),
**kwargs)
out = tf.nn.depthwise_conv2d(x, W, stride, padding)
if bias:
outdim = kernel[2] * multiplier
b = tf.get_bias('b',
shape=(outdim, ),
initializer=tf.zeros_initializer(),
**kwargs)
out = tf.nn.bias_add(out, b)
return out
def time_to_batch(data3d, dilation, paddings=None):
"""
:param data3d: [batch, time, channel]
:param dilation: int, rate or dilation
:param paddings: None or padding formats
:return: roughly [batch * dilation, time / dilation, channel] with padding if need
"""
# calc additional pad for depth dim
batch, width, channel = data3d.dims
# padding = dilation - (width - (width//dilation) * dilation)
# if padding != dilation:
# # if need, add padding
# data3d = tf.pad(data3d, [(0,0), (0, padding), (0, 0)])
if paddings is not None:
data3d = tf.pad(data3d, paddings)
# reshape, dim of time is shrinked along with dilation rate
data3d = data3d.reshape(-1, dilation, channel)
data3d = data3d.transpose(1, 0, 2)
data3d = data3d.reshape(batch * dilation, -1, channel)
return data3d
def conv3d(x,
outdim,
kernel,
stride=1,
pad=0,
padding='SAME',
mode='CONSTANT',
initializer=tf.he_uniform,
bias=False,
**kwargs):
kernel = _kernel_shape(3, kernel, x.dims[-1], outdim)
stride = _stride_shape(3, stride) # stride 5-dim
pads = None
if padding == 'SAME' and mode != 'CONSTANT':
# pad manually
half = ((kernel[0] - 1) // 2, (kernel[1] - 1) // 2,
(kernel[2] - 1) // 2)
pads = [(0, 0), (pad + half[0], pad + kernel[0] - 1 - half[0]),
(pad + half[1], pad + kernel[1] - 1 - half[1]),
(pad + half[2], pad + kernel[2] - 1 - half[2]), (0, 0)]
padding = 'VALID' # change to valid because manually padded
elif pad:
pads = [(0, 0), (pad, pad), (pad, pad), (pad, pad), (0, 0)]
if pads is not None:
x = tf.pad(x, pads, mode=mode)
W = tf.get_weight('W',
shape=kernel,
initializer=initializer(kernel),
**kwargs)
out = tf.nn.conv3d(x, W, stride, padding)
if bias:
b = tf.get_bias('b',
shape=(outdim, ),
initializer=tf.zeros_initializer(),
**kwargs)
out = tf.nn.bias_add(out, b)
return out
def pad_if_need(image, size, offsets=None):
"""
:param image: tensor3d[H,W,C]
:param size: (int, int) targetsize (H,W)
:param offsets: (0,0) for None
:return:
"""
assert image.ndim == 3
imshape = tf.shape(image)
# get target shape if possible
tshape = image.dims
for i in (0, 1):
if tshape[i] is not None and size[i] > tshape[i]:
tshape[i] = size[i]
targetshape = tf.convert_to_tensor(size).append(imshape[-1])
need = targetshape - imshape
# padding need
need = tf.where(need > 0, need, tf.zeros(tf.shape(need), dtype=tf.int32))
if offsets is None:
offsets = [0, 0, 0]
else:
offsets = list(offsets)
offsets.append(0)
# upper padding = need // 2
padding_first = need // 2 + tf.convert_to_tensor(offsets)
padding_left = need - padding_first
padding = tf.concat(0, [[padding_first], [padding_left]]).T
out = tf.pad(image, padding, 'CONSTANT')
# rshape = tf.maximum(imshape, targetshape)
# if known shape.. set
out.set_shape(tshape)
return out
def atrous(x,
outdim,
kernel,
rate,
pad=0,
padding='SAME',
initializer=tf.he_uniform,
bias=None,
**kwargs):
# todo rate per axis?
assert isinstance(pad, int)
nd = x.ndim - 2
if pad:
pads = [(0, 0)] + [(pad, pad)] * nd + [(0, 0)]
x = tf.pad(x, pads, mode='CONSTANT')
kernel = _kernel_shape(nd, kernel, x.dims[-1], outdim)
W = tf.get_weight('W',
shape=kernel,
initializer=initializer(kernel),
**kwargs)
if nd == 1:
out = _atrous1d(x, W, rate, padding=padding)
elif nd == 2:
out = tf.nn.atrous_conv2d(x, W, rate, padding)
else:
raise NotImplementedError('not implementd for ndim [{0}]'.format(nd))
if bias is not None:
b = tf.get_bias('b',
shape=(outdim, ),
initializer=tf.zeros_initializer(),
**kwargs)
out = tf.nn.bias_add(out, b)
return out