def __init__(self, input_shape: nc.TensorShape,
affine_shape: nc.TensorShape, output_size):
if input_shape.rank != 4:
raise ValueError('input_shape must be rank 4 (NCHW)')
N, IC, IH, IW = input_shape
if affine_shape.rank != 3:
raise ValueError('affine_shape must be rank 3')
AN, AH, AW = affine_shape
self.aff_tile = 0
if AN != N:
raise ValueError(f'affine matrix batch must == {AN}')
if AH != 2:
raise ValueError('affine matrix height must == 2')
if AW != 3:
raise ValueError('affine matrix width must == 3')
if output_size is not None:
OH, OW = output_size
else:
OH, OW = IH, IW
self.coords_shape = nc.TensorShape((N, IC, OH, OW, 3))
self.coords_init = nc.initializer.CoordsArange(0, IH - 1, 0, IW - 1)
self.coords_reshape = nc.TensorShape((N, IC * OH * OW, 3))
self.coords_affined_shape = nc.TensorShape((N, IC, OH, OW, 2))
self.output_shape = output_shape = nc.TensorShape((N, IC, OH, OW))
def __init__(self, shape, axes, keepdims):
shape_axes = shape.axes_arange()
if axes.is_none_axes():
axes = shape_axes
# Check correctness of axes
for axis in axes:
if axis not in shape_axes:
raise ValueError(f'Wrong axis {axis} not in {shape_axes}')
self.reduction_axes = reduction_axes = axes.sorted()
# Output axes. Remove axes from shape_axes
self.output_axes = output_axes = shape_axes - axes
if output_axes.is_none_axes():
output_shape = nc.TensorShape((1, ))
else:
output_shape = shape[output_axes]
self.output_shape = output_shape
self.output_shape_keepdims = nc.TensorShape(
1 if axis in reduction_axes else shape[axis]
for axis in range(shape.rank))
if keepdims:
self.output_shape = self.output_shape_keepdims
def __init__(self, t, shape):
shape = nc.TensorShape(shape)
if t.shape.size != shape.size:
raise ValueError(
f'Cannot interpet shape {t.shape} as ref shape {shape}')
super().__init__(shape)
self._t = t
def __init__(self, shape, tiles):
if len(tiles) != shape.rank:
raise ValueError(f'tiles should match shape.rank {shape.rank}')
self.output_shape = nc.TensorShape(dim * tiles[i]
for i, dim in enumerate(shape))
c = [0] * shape.rank
axes_offsets = []
for n in range(np.prod(tiles)):
axes_offsets.append(c.copy())
for i in range(shape.rank - 1, -1, -1):
c[i] += 1
if c[i] < tiles[i]:
break
c[i] = 0
axes_slices = []
for axes_offset in axes_offsets:
sl = []
for axis, tile in enumerate(axes_offset):
axis_size = shape[axis]
sl.append(slice(axis_size * tile, axis_size * (tile + 1)))
axes_slices.append(tuple(sl))
self.axes_slices = tuple(axes_slices)
def __init__(self, shape, init=None):
self.shape = shape = nc.TensorShape(shape)
if init is None:
init = nn.initializer.Scalar(0.0)
Tensor._object_count += 1
# ! internal variables (start with _) must be accesed through methods,
# ! because Tensor can be a TensorRef that uses method overridings
self._seq_id = Tensor._seq_id = Tensor._seq_id + 1
self._freezed = Tensor._freeze_stack != 0 # indicates that backprop
# should not go through inputs of _gradfns,
# which are marked as _is_trainable()
# use _is_freezed() to get the value
self._trainable = False # indicates that Tensor is used in Optimizer
# use _is_trainable() to get the value
self._grad_t = None # Gradient Tensor of this tensor
# use get_grad() to get the value
self._gradfns = None # Dict of [input_tensor] = func(O_t,dO_t), see _assign_grad
# use _get_gradfns() to get the value
self._op_name = None # Name of op which produces this Tensor
self._parent_module_w = None # weakref of parent Module
super().__init__(
shape.size * 4,
init_func=lambda x: init.initialize_CLBuffer(x, shape)
if init is not None else None)
def resize2D_bilinear(input_t, size_or_output_hw):
"""
resize2D_bilinear operator
arguments
size_or_output_hw int
float
Iterable of height,weight
"""
N, C, H, W = input_t.shape
if isinstance(size_or_output_hw, Iterable):
OH, OW = int(size_or_output_hw[0]), int(size_or_output_hw[1])
elif isinstance(size_or_output_hw, (int, float)):
OH = int(H * size_or_output_hw)
OW = int(W * size_or_output_hw)
else:
raise ValueError(
f'Unknown type of size_or_output_hw : {size_or_output_hw.__class__.__name__}'
)
OH = max(1, OH)
OW = max(1, OW)
coords_shape = nc.TensorShape((OH, OW, 2))
coords_t = nn.Tensor(coords_shape,
nn.initializer.CoordsArange(0, H - 1, 0, W - 1))
output_t = nn.spatial_transform2D(input_t, coords_t, grad_to_coords=False)
return output_t
def __init__(self, shape, axis, stack_count):
if axis < 0:
axis = shape.rank + 1 + axis
if axis < 0 or axis > shape.rank:
raise ValueError(f'Wrong axis {axis}')
if stack_count <= 0:
raise ValueError(f'Invalid stack_count {stack_count}')
self.output_shape = nc.TensorShape(
tuple(shape)[0:axis] + (stack_count, ) + tuple(shape)[axis:])
self.axis = axis
def copy(self, shape=None):
"""
Creates new tensor with the same shape and data.
shape override with new shape, but should match shape.size
"""
if shape is None:
shape = self.shape
else:
shape = nc.TensorShape(shape)
if shape.size != self.shape.size:
raise ValueError('shapes size mismatch')
t = Tensor(shape)
t.set(self)
return t
def __init__(self, input_shape: nc.TensorShape, size, is_add_to_output):
N, IC, IH, IW = input_shape
OC = IC
OH = IH * size
OW = IW * size
self.output_shape = output_shape = nc.TensorShape((N, OC, OH, OW))
common_kernel_text = f"""
{ph.define_axes_accessor('I', input_shape, 'NCHW')}
{ph.define_axes_accessor('O', output_shape, 'NCHW')}
"""
self.O_forward_krn = nc.CLKernel(global_shape=(output_shape.size, ),
kernel_text=f"""
{common_kernel_text}
#define SIZE {size}
__kernel void impl(__global float* O, __global const float* I)
{{
size_t gid = get_global_id(0);
{ph.axes_idxs_from_var('o', 'NCHW', 'gid')}
O[gid] {'+=' if is_add_to_output else '='} I[I_idx(on,oc,oh / SIZE,ow / SIZE)];
}}
""")
self.dI_krn = nc.CLKernel(global_shape=(input_shape.size, ),
kernel_text=f"""
{common_kernel_text}
#define SIZE {size}
__kernel void impl(__global float* dI, __global const float* dO)
{{
size_t gid = get_global_id(0);
{ph.axes_idxs_from_var('i', 'NCHW', 'gid')}
float v = 0.0;
for (int y=0; y<SIZE; ++y)
for (int x=0; x<SIZE; ++x)
v += dO[O_idx(in,ic,ih*SIZE+y,iw*SIZE+x)];
dI[gid] += v;
}}
""")
def tile_test():
for _ in range(10):
for shape_len in range(3, 5):
try:
shape = tuple(np.random.randint( 8, size=(shape_len,) )+1)
tiles = tuple(np.random.randint( 4, size=(shape_len,) )+1)
val_n = np.random.randint( 2**8, size=shape ).astype(np.float32)
tiled_n = np.tile(val_n, tiles)
val_t = nn.Tensor_from_value(val_n)
tiled_t = nn.tile(val_t, tiles)
if tiled_n.shape != tiled_t.shape:
raise Exception(f'shape is not equal')
if not all ( np.ndarray.flatten( tiled_t.np() == tiled_n ) ):
raise Exception(f'data is not equal')
tiled_n_grad = np.random.randint( 2**8, size=tiled_n.shape ).astype(np.float32)
val_t.get_grad().fill(1.0)
nn.backward( {tiled_t:tiled_n_grad} , grad_for_non_trainables=True )
info = nc.info.InfoTile( nc.TensorShape(shape), tiles)
val_n_grad = sum([ tiled_n_grad[axes_slice] for axes_slice in info.axes_slices ])
if not all ( np.ndarray.flatten(val_t.get_grad().np()-1.0 == val_n_grad) ):
raise Exception(f'grad is not equal')
except:
raise Exception(f"""
shape : {shape}
tiles : {tiles}
tiled_n_shape : {tiled_n.shape}
tiled_t_shape : {tiled_t.shape}
exception : {traceback.format_exc()}
""")
def __init__(self, shape, target_shape):
output_shape = []
remain_size = shape.size
unk_axis = None
for t_size in target_shape:
t_size = int(t_size)
if t_size != -1:
mod = remain_size % t_size
if mod != 0:
raise ValueError(
f'Cannot reshape {shape} to {target_shape}.')
remain_size /= t_size
else:
if unk_axis is not None:
raise ValueError('Can specify only one unknown dimension.')
unk_axis = len(output_shape)
output_shape.append(t_size)
if unk_axis is not None:
output_shape[unk_axis] = int(remain_size)
self.output_shape = nc.TensorShape(output_shape)
def __init__(self, shapes, axis):
shapes = tuple(shapes)
if len(shapes) == 0:
raise ValueError('shapes is empty')
shape = shapes[0]
if axis < 0:
axis = shape.rank + axis
if axis < 0 or axis >= shape.rank:
raise ValueError(f'Wrong axis {axis}')
fixed_shapes = [
tuple(a for i, a in enumerate(shape) if i != axis)
for shape in shapes
]
req_shape = fixed_shapes[0]
if not all(shape == req_shape for shape in fixed_shapes[1:]):
raise ValueError(
f'All shapes must match shape {tuple(a if i != axis else "*" for i,a in enumerate(shape))}'
)
axis_sizes = [shape[axis] for shape in shapes]
axis_offset = 0
axis_offsets = []
for axis_size in axis_sizes:
axis_offsets.append(axis_offset)
axis_offset += axis_size
self.output_shape = nc.TensorShape(
tuple(shape)[0:axis] + (sum(axis_sizes), ) +
tuple(shape)[axis + 1:])
self.axis = axis
self.axis_sizes = axis_sizes
self.axis_offsets = tuple(axis_offsets)
def __init__(self, input_shape: nc.TensorShape,
coords_shape: nc.TensorShape):
N, IC, IH, IW = input_shape
if coords_shape.rank not in [3, 4, 5]:
raise ValueError(
f'Coords shape rank must be 3(HWD) or 4(CHWD) or 5(NCHWD)')
KN, KC = 1, 1
if coords_shape.rank == 5:
KN, KC, KH, KW, KD = coords_shape
elif coords_shape.rank == 4:
KC, KH, KW, KD = coords_shape
elif coords_shape.rank == 3:
KH, KW, KD = coords_shape
self.coords_N_tile = 1
self.coords_C_tile = 1
if KN != N:
if KN == 1:
self.coords_N_tile = N
else:
raise ValueError(
f'Coords output batch {KN} does not match tensor input batch {N}.'
)
if KC != IC:
if KC == 1:
self.coords_C_tile = IC
else:
raise ValueError(
f'Coords output channels {KC} does not match tensor input channels {IC}.'
)
if KD != 2:
raise ValueError(f'Coords D {KD} channels must be == 2 (x,y)')
self.output_shape = output_shape = nc.TensorShape((N, IC, KH, KW))
common_kernel_text = f"""
{ph.define_axes_accessor('I', input_shape, 'NCHW')}
{ph.define_axes_accessor('O', output_shape, 'NCHW')}
"""
self.O_forward_krn = nc.CLKernel(global_shape=(output_shape.size, ),
kernel_text=f"""
{common_kernel_text}
{ph.define_axes_accessor('K', (KN,KC,KH,KW), 'NCHW')}
__kernel void impl(__global float* O, __global const float* I, __global const float2* K)
{{
size_t gid = get_global_id(0);
{ph.axes_idxs_from_var('o', 'NCHW', 'gid')}
float v = 0.0;
float2 xys = K[K_idx_mod(on,oc,oh,ow)];
for (int ih=0; ih < IH; ++ih)
{{
float ys_mod = max(0.0, 1.0-fabs(xys.y-ih));
if (ys_mod != 0.0)
for (int iw=0; iw < IW; ++iw)
{{
float xs_mod = max(0.0, 1.0-fabs(xys.x-iw));
if (xs_mod != 0.0)
v += xs_mod*ys_mod*I[I_idx(on,oc,ih,iw)];
}}
}}
O[gid] = v;
}}
""")
self.dI_krn = nc.CLKernel(global_shape=(input_shape.size, ),
kernel_text=f"""
{common_kernel_text}
{ph.define_axes_accessor('K', (KN,KC,KH,KW), 'NCHW')}
__kernel void impl(__global float* dI, __global const float2* K, __global float* dO)
{{
size_t gid = get_global_id(0);
{ph.axes_idxs_from_var('i', 'NCHW', 'gid')}
float v = 0.0;
for (int oh=0; oh < OH; ++oh)
for (int ow=0; ow < OW; ++ow)
{{
float2 xys = K[K_idx_mod(in,ic,oh,ow)];
float ys_mod = max(0.0, 1.0-fabs(xys.y-ih));
if (ys_mod != 0.0)
{{
float xs_mod = max(0.0, 1.0-fabs(xys.x-iw));
if (xs_mod != 0.0)
v += xs_mod*ys_mod*dO[O_idx(in,ic,oh,ow)];
}}
}}
dI[gid] += v;
}}
""")
self.dK_krn = nc.CLKernel(global_shape=(N * IC * KH * KW, ),
kernel_text=f"""
{common_kernel_text}
{ph.define_axes_accessor('K', (N,IC,KH,KW), 'NCHW')}
__kernel void impl(__global float2* dK, __global const float* I, __global const float2* K, __global float* dO)
{{
size_t gid = get_global_id(0);
{ph.axes_idxs_from_var('k', 'NCHW', 'gid')}
float dk_x = 0.0;
float dk_y = 0.0;
float2 xys = K[gid];
for (int ih=0; ih < IH; ++ih)
{{
{{
float ys_mod = max(0.0, 1.0-fabs(xys.y-ih));
if (ys_mod != 0.0)
for (int iw=0; iw < IW; ++iw)
if ( fabs(iw-xys.x) < 1.0 )
{{
float xs_mod = 1.0*(iw >= xys.x)-1.0*(iw < xys.x);
dk_x += xs_mod*ys_mod*I[I_idx(kn,kc,ih,iw)] * dO[O_idx(kn,kc,kh,kw)];
}}
}}
if (fabs(ih-xys.y) < 1.0)
{{
float ys_mod = 1.0*(ih >= xys.y)-1.0*(ih < xys.y);
for (int iw=0; iw < IW; ++iw)
{{
float xs_mod = max(0.0, 1.0-fabs(xys.x-iw));
if (xs_mod != 0.0)
dk_y += xs_mod*ys_mod*I[I_idx(kn,kc,ih,iw)] * dO[O_idx(kn,kc,kh,kw)];
}}
}}
}}
dK[gid] += (float2)(dk_x, dk_y);
}}
""")
def __init__(self, op_type, input_shape: nc.TensorShape, pool_size, stride,
padding):
if op_type not in ['avg', 'min', 'max']:
raise ValueError(f'unknown op_type {op_type}')
if pool_size < 2:
raise ValueError(f'pool_size {pool_size} must be at least 2')
self.op_type = op_type
N, IC, IH, IW = input_shape
ci = nc.info.InfoConv2D(IH, IW, pool_size, pool_size, stride, 1,
padding)
OC, OH, OW = IC, ci.OH, ci.OW
self.output_shape = output_shape = nc.TensorShape((N, OC, OH, OW))
common_kernel_text = f"""
{ph.define_axes_accessor('I', input_shape, 'NCHW')}
{ph.define_axes_accessor('O', output_shape, 'NCHW')}
#define PADL {ci.PADL}
#define PADT {ci.PADT}
#define POOL_SIZE {pool_size}
#define STRIDE {stride}
"""
self.O_forward_krn = nc.CLKernel(global_shape=(output_shape.size, ),
kernel_text=f"""
{common_kernel_text}
__kernel void impl(__global float* O, __global const float* I)
{{
size_t gid = get_global_id(0);
{ph.axes_idxs_from_var('o', 'NCHW', 'gid')}
{ {'avg' : 'float v = 0.0; int v_count = 0;',
'max' : 'float v = -INFINITY;',
'min' : 'float v = INFINITY;',
}[op_type] }
for (int ph=0; ph<POOL_SIZE; ++ph)
for (int pw=0; pw<POOL_SIZE; ++pw)
{{
int ih = -PADT + ph + oh*STRIDE;
int iw = -PADL + pw + ow*STRIDE;
if (iw >= 0 & ih >= 0 & iw < IW & ih < IH)
{{
{ {'avg' : 'v += I[I_idx(on,oc,ih,iw)]; ++v_count;',
'max' : 'v = fmax(v, I[I_idx(on,oc,ih,iw)]);',
'min' : 'v = fmin(v, I[I_idx(on,oc,ih,iw)]);',
}[op_type] }
}}
}}
{ {'avg' : 'if (v_count != 0) v /= v_count;',
'max' : 'if (v == -INFINITY) v = 0.0;',
'min' : 'if (v == INFINITY) v = 0.0;',
}[op_type] }
O[gid] = v;
}}
""")
if op_type == 'avg':
self.dI_krn = nc.CLKernel(global_shape=(input_shape.size, ),
kernel_text=f"""
{common_kernel_text}
__kernel void impl(__global float* dI, __global const float* dO)
{{
size_t gid = get_global_id(0);
{ph.axes_idxs_from_var('i', 'NCHW', 'gid')}
float v = 0.0;
for (int ph=0; ph<POOL_SIZE; ++ph)
for (int pw=0; pw<POOL_SIZE; ++pw)
{{
int oh = (PADT + ih - ph ) / STRIDE;
int ow = (PADL + iw - pw ) / STRIDE;
if (ow >= 0 & oh >= 0 & ow < OW & oh < OH
& iw == (-PADL + pw + ow*STRIDE)
& ih == (-PADT + ph + oh*STRIDE) )
{{
int d=0;
for (int dph=0; dph<POOL_SIZE; ++dph)
for (int dpw=0; dpw<POOL_SIZE; ++dpw)
{{
int dih = -PADT + dph + oh*STRIDE;
int diw = -PADL + dpw + ow*STRIDE;
d += (diw >= 0 & dih >= 0 & diw < IW & dih < IH);
}}
v += dO[O_idx(in,ic,oh,ow)] / d;
}}
}}
dI[gid] += v;
}}
""")
elif op_type in ['min', 'max']:
# Implementation is different from tensorflow in case when the same values exist in reduction axes.
# Example tf : 3 4 5 5 , max = 5, gradients : 0 0 0.5 0.5
# Example litenn : 3 4 5 5 , max = 5, gradients : 0 0 1 0
# or gradients : 0 0 0 1 - depends on which GPU thread will be first !
self.dI_krn = nc.CLKernel(global_shape=(input_shape.size, ),
kernel_text=f"""
{common_kernel_text}
__kernel void impl(__global float* dI, __global const float* I, __global const float* dO, __global const float* O)
{{
size_t gid = get_global_id(0);
{ph.axes_idxs_from_var('i', 'NCHW', 'gid')}
float v = 0.0;
// Iterate over all O pixels, where 'I' have contribution
for (int ph=0; ph<POOL_SIZE; ++ph)
for (int pw=0; pw<POOL_SIZE; ++pw)
{{
int oh = (PADT + ih - ph ) / STRIDE;
int ow = (PADL + iw - pw ) / STRIDE;
if (ow >= 0 & oh >= 0 & ow < OW & oh < OH
& iw == (-PADL + pw + ow*STRIDE)
& ih == (-PADT + ph + oh*STRIDE) )
{{
//Now we have oh,ow where ih,iw have contribution
float Ov = O[O_idx(in,ic,oh,ow)];
//Iterate in those I pixels, which were used to produce O
//to determine first min/max match
for (int dphw=0; dphw < POOL_SIZE*POOL_SIZE; ++dphw)
{{
int dih = -PADT + (dphw / POOL_SIZE) + oh*STRIDE;
int diw = -PADL + (dphw % POOL_SIZE) + ow*STRIDE;
if (diw >= 0 & dih >= 0 & diw < IW & dih < IH &
I[I_idx(in,ic,dih,diw)] == Ov)
{{
// Match I==O
if (dih == ih & diw == iw)
// but add gradient only if current ih/iw index match dih/diw
v += dO[O_idx(in,ic,oh,ow)];
break;
}}
}}
}}
}}
dI[gid] += v;
}}
""")
def __init__(self, input_shape : nc.TensorShape, kernel_shape : nc.TensorShape, stride, dilation, padding):
if kernel_shape.rank != 3:
raise ValueError(f'Kernel shape rank must be == 3')
N,IC,IH,IW = input_shape
KI,KH,KW = kernel_shape
if KI != IC:
raise ValueError(f'Kernel input channels {KI} does not match tensor input channels {IC}.')
ci = nc.info.InfoConv2D(IH, IW, KH, KW, stride, dilation, padding)
OC, OH, OW = IC, ci.OH, ci.OW
self.output_shape = output_shape = nc.TensorShape( (N, OC, OH, OW) )
self.OC_1_1_NxOHxOW = (OC,1,1,N*OH*OW)
self.KI_KH_KW_NxOHxOW = (KI,KH,KW,N*OH*OW)
common_kernel_text = f"""
{ph.define_axes_accessor('I', input_shape, 'NCHW')}
{ph.define_axes_accessor('O', output_shape, 'NCHW')}
{ph.define_axes_accessor('K', kernel_shape, 'IHW')}
#define PADL {ci.PADL}
#define PADT {ci.PADT}
#define STRIDE {stride}
#define DILATION {dilation}
"""
self.O_depthwise_krn = nc.CLKernel(global_shape=(output_shape.size,), kernel_text=f"""
{common_kernel_text}
__kernel void impl(__global float* O, __global const float* I, __global const float* K)
{{
size_t gid = get_global_id(0);
{ph.axes_idxs_from_var('o', 'NCHW', 'gid')}
float v = 0.0;
for (int kh=0; kh<KH; ++kh)
{{
int ih = -PADT + kh*DILATION + oh*STRIDE;
if (ih >= 0 & ih < IH)
for (int kw=0; kw<KW; ++kw)
{{
int iw = -PADL + kw*DILATION + ow*STRIDE;
if (iw >= 0 & iw < IW)
v += I[I_idx(on,oc,ih,iw)]*K[K_idx(oc,kh,kw)];
}}
}}
O[gid] = v;
}}
""")
self.dI_depthwise_krn = nc.CLKernel(global_shape=(input_shape.size,), kernel_text=f"""
{common_kernel_text}
__kernel void impl(__global float* dI, __global const float* K, __global const float* dO)
{{
size_t gid = get_global_id(0);
{ph.axes_idxs_from_var('i', 'NCHW', 'gid')}
float v = 0.0;
for (int kh=0; kh<KH; ++kh)
{{
int oh = (PADT + ih - kh*DILATION ) / STRIDE;
if (oh >= 0 & oh < OH)
for (int kw=0; kw<KW; ++kw)
{{
int ow = (PADL + iw - kw*DILATION ) / STRIDE;
if (ow >= 0 & ow < OW
& iw == (-PADL + kw*DILATION + ow*STRIDE)
& ih == (-PADT + kh*DILATION + oh*STRIDE) )
v += dO[O_idx(in,ic,oh,ow)]*K[K_idx(ic,kh,kw)];
}}
}}
dI[gid] += v;
}}
""")
self.im2col = lambda x: nc.op.unfold2D(x, N, IC, IH, IW, OH, OW, KH, KW, ci.PADL, ci.PADT, dilation, stride, 'CJI_NHW', is_transpose=False)
def __init__(self, input_shape: nc.TensorShape, slices, is_add_to_output):
# Validate slices argument for given shape.
new_slices = []
before_ellipsis = None
for s in slices:
if s is Ellipsis:
before_ellipsis = new_slices
new_slices = []
continue
elif s is not None and not isinstance(s, (int, tuple)):
raise ValueError(
f'unknown slice argument {s} of type {s.__class__}')
new_slices.append(s)
if before_ellipsis is not None:
# Process Ellipsis separator
new_slices_n_axes = sum([1 for x in new_slices if x != None])
before_ellipsis_n_axes = sum(
[1 for x in before_ellipsis if x != None])
# Expand slices by filling intermediate (None,None,None) for each remaining axis
new_slices = before_ellipsis + \
[(None,None,None)]*max(0, input_shape.rank-before_ellipsis_n_axes-new_slices_n_axes) + \
new_slices
new_slices_n_axes = sum([1 for x in new_slices if x != None])
if new_slices_n_axes > input_shape.rank:
raise ValueError('slices arguments more than shape axes')
elif new_slices_n_axes < input_shape.rank:
# Fill remaining axes
new_slices += [(None, None, None)
] * (input_shape.rank - new_slices_n_axes)
slices = tuple(new_slices)
# Compute shapes
output_is_reshaped = True # Flag determines that output_tensor
# can be just reshaped without any computation
output_shape = [] # output tensor shape
output_shape_krn = [
] # output shape used in kernel, must match input shape
input_axes_begin_step = [] # begin,step ints for every input shape
i_axis = 0
for v in slices:
if v is None:
# None is new axis
# We can add unlimited number of (1,) axes at any place of shape
output_shape.append(1)
continue
i_axis_size = input_shape[i_axis]
i_axis += 1
if isinstance(v, int):
if v < 0: v += i_axis_size
if v < 0 or v >= i_axis_size:
raise ValueError(
f'index {v} is out of bounds for axis {i_axis} with size {i_axis_size}'
)
b, e, s = v, v, 1
else:
b, e, s = v
# Fix begin, end, step values
if s is None: s = 1
if s == 0:
raise ValueError('slice step cannot be zero')
if b is None: b = 0 if s > 0 else i_axis_size - 1
if e is None: e = i_axis_size if s > 0 else -1
elif e < 0: e += i_axis_size
if b < 0: b += i_axis_size
if s > 0:
b = np.clip(b, 0, i_axis_size)
e = np.clip(e, 0, i_axis_size)
else:
b = np.clip(b, 0, i_axis_size - 1)
e = np.clip(e, -1, i_axis_size)
if i_axis_size != 1 and not (b == 0 and e == i_axis_size
and s == 1):
# Such params of axis slice will change input, thus output cannot be just reshaped input
output_is_reshaped = False
# Compute output_axis_size based on begin,end,step
output_axis_size = max(0, math.ceil((e - b) / s))
if output_axis_size >= 1:
# >= 1 : select range of indexes, axis will remain
output_shape.append(output_axis_size)
# ^ othwerwise axis will be supressed
# output_shape to use in kernel, must match rank of input shape
output_shape_krn.append(max(1, output_axis_size))
# for every output_shape_krn axis
# we have exact begin,step values to fetch value from input
input_axes_begin_step.append((b, s))
output_shape_krn = nc.TensorShape(output_shape_krn)
self.output_is_reshaped = output_is_reshaped
self.output_shape = nc.TensorShape(output_shape)
self.forward_krn = nc.CLKernel(global_shape=(output_shape_krn.size, ),
kernel_text=f"""
{ph.define_axes_accessor('I', input_shape )}
{ph.define_axes_sizes('O', output_shape_krn )}
__kernel void impl(__global const float* I, __global float* O)
{{
size_t gid = get_global_id(0);
{ph.axes_idxs_from_var('o', output_shape_krn.rank, 'gid')}
{''.join( f'size_t i{i} = {b} + o{i} * {s};' for i, (b,s) in enumerate(input_axes_begin_step) ) }
O[get_global_id(0)] {'+=' if is_add_to_output else '='} I[I_idx({ph.axes_seq_enum('i', input_shape.rank)})];
}}
""")
self.backward_krn = nc.CLKernel(global_shape=(output_shape_krn.size, ),
kernel_text=f"""
{ph.define_axes_accessor('I', input_shape )}
{ph.define_axes_sizes('O', output_shape_krn )}
__kernel void impl(__global float* dI, __global const float* O)
{{
size_t gid = get_global_id(0);
{ph.axes_idxs_from_var('o', output_shape_krn.rank, 'gid')}
{''.join( f'size_t i{i} = {b} + o{i} * {s};' for i, (b,s) in enumerate(input_axes_begin_step) ) }
dI[I_idx({ph.axes_seq_enum('i', input_shape.rank)})] += O[get_global_id(0)];
}}
""")
def __init__(self, input_shape, kernel_shape, stride, dilation, padding):
if padding not in ['valid', 'same']:
raise ValueError(
'Wrong padding value, only valid or same supported for conv2DTranspose.'
)
N, IC, IH, IW = input_shape
KO, KI, KH, KW = kernel_shape
if KI != IC:
raise ValueError(
f'Kernel input channels {KI} does not match input channels {IC}.'
)
ci = nc.info.InfoConv2D(IH, IW, KH, KW, stride, dilation, padding)
OC, OH, OW = KO, ci.OH_T, ci.OW_T
self.output_shape = output_shape = nc.TensorShape((N, OC, OH, OW))
self.OC_N_OH_OW = (OC, N, OH, OW)
self.OC_NxOHxOW = (OC, N * OH * OW)
self.N_IH_IW_IC = (N, IH, IW, IC)
self.im2colT = lambda x: nc.op.unfold2D(x,
N,
IC,
IH,
IW,
OH,
OW,
KH,
KW,
ci.PADL,
ci.PADT,
dilation,
stride,
'CJI_NHW',
is_transpose=True)
self.im2rowT = lambda x: nc.op.unfold2D(x,
N,
IC,
IH,
IW,
OH,
OW,
KH,
KW,
ci.PADL,
ci.PADT,
dilation,
stride,
'NHW_CJI',
is_transpose=True)
self.im2row = lambda x: nc.op.unfold2D(x,
N,
OC,
OH,
OW,
IH,
IW,
KH,
KW,
ci.PADL,
ci.PADT,
dilation,
stride,
'NHW_CJI',
is_transpose=False)
def __init__(self, input_shape, kernel_shape, stride, dilation, padding):
N, IC, IH, IW = input_shape
KO, KI, KH, KW = kernel_shape
if KI != IC:
raise ValueError(
f'Kernel input channels {KI} does not match tensor input channels {IC}.'
)
ci = nc.info.InfoConv2D(IH, IW, KH, KW, stride, dilation, padding)
OC, OH, OW = KO, ci.OH, ci.OW
self.output_shape = output_shape = nc.TensorShape((N, OC, OH, OW))
self.OC_N_OH_OW = (OC, N, OH, OW)
self.OC_NxOHxOW = (OC, N * OH * OW)
self.N_IH_IW_IC = (N, IH, IW, IC)
self.im2col = lambda x: nc.op.unfold2D(x,
N,
IC,
IH,
IW,
OH,
OW,
KH,
KW,
ci.PADL,
ci.PADT,
dilation,
stride,
'CJI_NHW',
is_transpose=False)
self.im2row = lambda x: nc.op.unfold2D(x,
N,
IC,
IH,
IW,
OH,
OW,
KH,
KW,
ci.PADL,
ci.PADT,
dilation,
stride,
'NHW_CJI',
is_transpose=False)
self.im2rowT = lambda x: nc.op.unfold2D(x,
N,
OC,
OH,
OW,
IH,
IW,
KH,
KW,
ci.PADL,
ci.PADT,
dilation,
stride,
'NHW_CJI',
is_transpose=True)
def __init__(self, a_shape, b_shape, is_add_to_output):
if a_shape.rank != b_shape.rank:
raise ValueError(
f'Ranks are not equal. {a_shape.rank} != {b_shape.rank}')
rank = a_shape.rank
if rank < 2:
raise ValueError('Tensors rank must be at least 2.')
K, M = a_shape[-2], a_shape[-1]
N, B_COLS = b_shape[-2], b_shape[-1]
if K != B_COLS:
raise ValueError('A_ROWS != B_COLS')
BATCH = a_shape[0:-2].size
B_BATCH = b_shape[0:-2].size
if BATCH != B_BATCH:
raise ValueError(
f'BATCH size {BATCH} != {B_BATCH} in shapes {a_shape} {b_shape}'
)
if rank == 2:
self.output_shape = output_shape = nc.TensorShape((N, M))
else:
self.output_shape = output_shape = nc.TensorShape(a_shape[:-2] +
(N, M))
self.M = M
self.N = N
self.K = K
# Determining optimal tile widths
for MW in [16, 8, 4, 2, 1]:
if M % MW == 0:
break
for KW in [8, 4, 2, 1]:
if N % KW == 0 and K % KW == 0:
break
NW = KW
self.forward_krn = nc.CLKernel(global_shape=(M // MW, N // NW, BATCH),
kernel_text=f"""
#define K {K}
#define N {N}
#define MW {MW} // M tile Width
#define NW {NW} // N tile Width -- NW & KW should be the same !
#define KW {KW} // K tile Width
#define MT {M//MW} // MT is max for 'mt' (M tile count)
#define KT {K//KW} // KT is max for 'kt' (K tile count)
#define floatMW { f'float{MW}' if MW != 1 else 'float'}
#define floatKW { f'float{KW}' if KW != 1 else 'float'}
__kernel void GeMM(const __global floatMW* restrict A, const __global floatKW* restrict B, __global floatMW* C)
{{
size_t mt = get_global_id(0); //global M-tile id
size_t nc = get_global_id(1); //global N-tile id
size_t batch = get_global_id(2);
float AT[KW][MW]; // sub tiles
float BT[NW][KW];
float CT[NW][MW];
#pragma unroll
for (uint i=0; i<NW*MW; ++i) // zero CT tile
((float*) CT)[i] = 0.0;
for (uint kt=0; kt<KT; ++kt) // iterate over K-dim tiles
{{
#pragma unroll
for (uint k=0; k<KW; ++k) // every k-element inside K-dim tile
*( (floatMW*) AT[k] ) = A[batch*K*MT + (kt*KW + k)*MT + mt]; // store M-Width floats
#pragma unroll
for (uint n=0; n<NW; ++n) // every n-element inside N-dim tile
*( (floatKW*) BT[n] ) = B[batch*N*KT + (nc*NW + n)*KT + kt]; // store K-Width floats
#pragma unroll
for (uint k=0; k<KW; ++k)
#pragma unroll
for (uint n=0; n<NW; ++n) // sub tiles multiplication
#pragma unroll
for (uint m=0; m<MW; ++m)
CT[n][m] += AT[k][m] * BT[n][k];
}}
#pragma unroll
for (uint n=0; n<NW; ++n)
C[ batch*N*MT + (nc*NW + n)*MT + mt] {'+=' if is_add_to_output else '='}
*( (floatMW*) CT[n]);
}}""")
def dual_wise_op_test():
for op in [
add, binary_crossentropy, categorical_crossentropy, sub, max, min,
mul, div
]:
print(f'{op.__name__}()')
for _ in range(10):
if op == categorical_crossentropy:
shape_gen = [2]
else:
shape_gen = range(1, 5)
for shape_len in shape_gen:
try:
a_shape = tuple(
np.random.randint(8, size=(shape_len, )) + 1)
if op == categorical_crossentropy:
b_shape = a_shape
else:
if np.random.randint(2) == 0:
b_shape = tuple(
a_shape[np.random.randint(len(a_shape)):])
b_shape = (1, ) if len(b_shape) == 0 else b_shape
else:
b_shape = list(a_shape)
b_shape[np.random.randint(len(b_shape))] = 1
b_shape = tuple(b_shape)
shapes = [a_shape, b_shape]
if np.random.randint(2) == 0:
shapes = shapes[::-1]
a_shape, b_shape = shapes
a_n = np.random.randint(1, 2**8,
size=a_shape).astype(np.float32)
b_n = np.random.randint(1, 2**8,
size=b_shape).astype(np.float32)
a_t = nn.Tensor_from_value(a_n)
b_t = nn.Tensor_from_value(b_n)
r_t = op(a_t, b_t)
r_n_grad = np.random.randint(2**8, size=r_t.shape).astype(
np.float32)
a_t.get_grad().fill(1.0)
b_t.get_grad().fill(1.0)
nn.backward({r_t: r_n_grad}, grad_for_non_trainables=True)
if op == div:
# Test validness and gradient only for div
r_n = a_n / b_n
if r_n.shape != r_t.shape:
raise Exception(f'shapes are not equal')
if np.abs(np.sum(
(np.ndarray.flatten(r_t.np() - r_n)))) > 1:
raise Exception(f'data is not equal')
info = nc.info.InfoBroadcast(nc.TensorShape(a_shape),
nc.TensorShape(b_shape))
a_n_grad = r_n_grad / b_n
axes = info.a_shape_reduction_axes
if axes.rank == 0:
a_n_grad = a_n_grad.reshape(a_n.shape)
else:
a_n_grad = a_n_grad.sum(tuple(axes), keepdims=True)
b_n_grad = r_n_grad * (-a_n / (b_n**2))
axes = info.b_shape_reduction_axes
if axes.rank == 0:
b_n_grad = b_n_grad.reshape(b_n.shape)
else:
b_n_grad = b_n_grad.sum(tuple(axes), keepdims=True)
if np.abs(
np.sum(
(np.ndarray.flatten(a_t.get_grad().np() -
1.0 - a_n_grad)))) > 1:
raise Exception(f'grad A is not equal')
if np.abs(
np.sum(
(np.ndarray.flatten(b_t.get_grad().np() -
1.0 - b_n_grad)))) > 1:
raise Exception(f'grad B is not equal')
else:
if not a_t.has_grad():
raise Exception(f'a_t has no grad')
if not b_t.has_grad():
raise Exception(f'b_t has no grad')
except:
raise Exception(f"""
op : {op}
a_shape : {a_shape}
b_shape : {b_shape}
r_n_shape : {r_n.shape}
exception : {traceback.format_exc() }
""")