def matvec(mat, vec, axis=0, out=None, alpha=1.0, beta=0.0, allocator=memPool):
assert vec.dtype == mat.dtype and (mat.dtype == np.float32
or mat.dtype == np.float16)
assert vec.ndim == mat.ndim - 1 and 0 <= axis < 2
h, w = mat.shape[-2:]
if axis == 1:
assert mat.dimAt(-1) == vec.dimAt(-1)
block = (warpSize, 1, 1)
grid = (h, 1, prod(mat.shape[:-2]))
if out is None:
out = GPUArray.zeros(mat.shape[:-1],
dtype=mat.dtype,
allocator=allocator)
else:
assert out.shape == mat.shape[:-1]
fn = mulmod.vecMulOnRow if mat.dtype == np.float32 else mulmod.vecMulOnRowFP16
fn(out,
mat,
vec,
np.int32(w),
np.int32(h),
np.float32(alpha),
np.float32(beta),
block=block,
grid=grid)
else:
block = (NT, 1, 1)
grid = (roundUpDiv(w, block[0]), 1, prod(mat.shape[:-2]))
if out is None:
out = GPUArray.zeros(mat.shape[:-2] + (w, ),
dtype=mat.dtype,
allocator=allocator)
else:
assert out.shape == mat.shape[:-2] + (w, )
fn = mulmod.vecMulOnCol if mat.dtype == np.float32 else mulmod.vecMulOnColFP16
fn(out,
mat,
vec,
np.int32(w),
np.int32(h),
np.float32(alpha),
np.float32(beta),
block=block,
grid=grid)
return out
def reflectpadBackward(grad, pad, allocator=memPool):
if grad.ndim == 3:
batchsize, maps, outsize = grad.shape
lpad, rpad = pad
block = (warpSize, 1, 1)
grid = (roundUpDiv(outsize, warpSize), maps, batchsize)
insize = outsize - lpad - rpad
ingrad = GPUArray.zeros((batchsize, maps, insize),
dtype=grad.dtype,
allocator=allocator)
fn = mod.reflectpad1dBackward if grad.dtype == np.float32 else mod.reflectpad1dBackwardFP16
fn(ingrad,
grad,
np.int32(insize),
np.int32(lpad),
np.int32(rpad),
block=block,
grid=grid)
elif grad.ndim == 4:
batchsize, maps, outh, outw = grad.shape
upad, bpad, lpad, rpad = pad
inh, inw = outh - upad - bpad, outw - lpad - rpad
block = (warpSize, 1, 1)
grid = (roundUpDiv(outh * outw, warpSize), maps, batchsize)
ingrad = GPUArray.zeros((batchsize, maps, inh, inw),
dtype=grad.dtype,
allocator=allocator)
fn = mod.reflectpad2dBackward if grad.dtype == np.float32 else mod.reflectpad2dBackwardFP16
fn(ingrad,
grad,
np.int32(inh),
np.int32(inw),
np.int32(upad),
np.int32(bpad),
np.int32(lpad),
np.int32(rpad),
block=block,
grid=grid)
else:
raise NotImplementedError(grad.ndim)
return ingrad
def matsum(tensor, axis=0, out=None, alpha=1.0, beta=0.0, allocator=memPool):
assert tensor.dtype == np.float32 or tensor.dtype == np.float16
assert 0 <= axis < tensor.ndim
if axis == tensor.ndim - 1:
block = (warpSize, 1, 1)
grid = (prod(tensor.shape[:-1]), 1, 1)
if out is None:
out = GPUArray.zeros(tensor.shape[:-1],
dtype=tensor.dtype,
allocator=allocator)
else:
assert out.shape == tensor.shape[:-1]
fn = summod.sumOnRow if tensor.dtype == np.float32 else summod.sumOnRowFP16
fn(out,
tensor,
np.int32(tensor.dimAt(-1)),
np.float32(alpha),
np.float32(beta),
block=block,
grid=grid)
else:
z, width = prod(tensor.shape[:axis]), prod(tensor.shape[axis + 1:])
block = (NT, 1, 1)
grid = (roundUpDiv(width, block[0]), 1, z)
if out is None:
out = GPUArray.zeros(tensor.shape[:axis] + tensor.shape[axis + 1:],
dtype=tensor.dtype,
allocator=allocator)
else:
assert out.shape == tensor.shape[:axis] + tensor.shape[axis + 1:]
fn = summod.sumOnCol if tensor.dtype == np.float32 else summod.sumOnColFP16
fn(out,
tensor,
np.int32(width),
np.int32(tensor.dimAt(axis)),
np.float32(alpha),
np.float32(beta),
block=block,
grid=grid)
return out
def maxunpool2d(data, origshape, mask, allocator=memPool):
assert data.dtype == np.float32
batchsize, maps, inh, inw = data.shape
outh, outw = origshape[2], origshape[3]
outdata = GPUArray.zeros((batchsize, maps, outh, outw),
dtype=np.float32,
allocator=allocator)
size = prod(data.shape)
block = (nthreads, 1, 1)
grid = (roundUpDiv(size, nthreads), 1, 1)
mod.maxunpool2d(outdata,
data,
mask,
np.int32(inh),
np.int32(inw),
np.int32(outh),
np.int32(outw),
np.int32(maps),
np.int32(size),
block=block,
grid=grid)
return outdata
def upsample2dBackward(grad, scale, mode="nearest", allocator=memPool): batchsize, maps, outh, outw = grad.shape hscale, wscale = (scale, scale) if isinstance(scale, int) else scale inh, inw = outh // hscale, outw // wscale if mode == "nearest": ingrad = GPUArray.empty((batchsize, maps, inh, inw), dtype=grad.dtype, allocator=allocator) blk = warpSize * 8 block = (blk, 1, 1) grid = (roundUpDiv(ingrad.size, blk), 1, 1) nearestMod.upsample2dNearestBackward( ingrad, grad, np.int32(inw), np.int32(outw), np.int32(hscale), np.int32(wscale), np.int32(ingrad.size), block=block, grid=grid ) elif mode == "linear": ingrad = GPUArray.zeros((batchsize, maps, inh, inw), dtype=grad.dtype, allocator=allocator) block = (warpSize, nthreads // warpSize, 1) grid = (roundUpDiv(outw, block[0]), roundUpDiv(outh, block[1]), 1) rh, rw = (inh - 1) / (outh - 1), (inw - 1) / (outw - 1) linearMod.upsample2dLinearBackward( ingrad, grad, np.int32(batchsize), np.int32(maps), np.int32(inh), np.int32(inw), np.int32(outh), np.int32(outw), np.float32(rh), np.float32(rw), block=block, grid=grid ) else: raise NotImplementedError(mode) return ingrad
def batchNorm3dTest(dtype, atol): batchsize, maps, d, h, w = 2, 5, 2, 3, 2 epsilon, norm = 1e-5, batchsize * d * h * w hostData = np.random.randn(batchsize, maps, d, h, w).astype(dtype) hostScale = np.random.randn(1, maps, 1, 1, 1).astype(np.float32) hostBias = np.random.randn(1, maps, 1, 1, 1).astype(np.float32) data, scale, bias = GPUArray.toGpu(hostData), GPUArray.toGpu(hostScale.ravel()), GPUArray.toGpu(hostBias.ravel()) mean, var = GPUArray.zeros(scale.shape, dtype=np.float32), GPUArray.toGpu(np.ones(scale.shape, dtype=np.float32)) outdata, savemean, saveinvvar = context.batchNormNd(data, mean, var, scale, bias, epsilon=epsilon, out=data) hostMean = np.sum(hostData, axis=(0, 2, 3, 4), dtype=np.float32, keepdims=True) / norm hostInvVar = np.sum((hostData - hostMean) ** 2, axis=(0, 2, 3, 4), dtype=np.float32, keepdims=True) / norm hostInvVar = 1.0 / np.sqrt(hostInvVar + epsilon) hostNormData = (hostData - hostMean) * hostInvVar hostOutData = (hostNormData * hostScale + hostBias).astype(dtype) assert np.allclose(hostMean.ravel(), mean.get(), atol=atol) assert np.allclose(hostInvVar.ravel(), saveinvvar.get(), atol=atol) assert np.allclose(hostOutData, outdata.get(), atol=atol) hostGrad = np.random.randn(*outdata.shape).astype(dtype) grad, data = GPUArray.toGpu(hostGrad), GPUArray.toGpu(hostData) ingrad, scalegrad, biasgrad = context.batchNormNdBackward(grad, data, scale, savemean, saveinvvar, epsilon=epsilon) hostScaleGrad = np.sum(hostGrad * hostNormData, axis=(0, 2, 3, 4), dtype=np.float32, keepdims=True) hostBiasGrad = np.sum(hostGrad, axis=(0, 2, 3, 4), dtype=np.float32, keepdims=True) hostMeanGrad = -hostInvVar * hostBiasGrad * hostScale hostVarGrad = np.sum(hostGrad * (hostData - hostMean), axis=(0, 2, 3, 4), dtype=np.float32, keepdims=True) hostVarGrad = -0.5 * hostVarGrad * hostScale * hostInvVar**3 hostInGrad = hostGrad * hostScale * hostInvVar + (2 * hostVarGrad * (hostData - hostMean) + hostMeanGrad) / norm hostInGrad = hostInGrad.astype(dtype) assert np.allclose(hostInGrad, ingrad.get(), atol=atol) assert np.allclose(hostScaleGrad.ravel(), scalegrad.get(), atol=atol) assert np.allclose(hostBiasGrad.ravel(), biasgrad.get(), atol=atol) hostMean = np.random.randn(*hostMean.shape).astype(np.float32) hostVar = 1.0 + np.random.randn(*hostInvVar.shape).astype(np.float32)**2 mean, var = GPUArray.toGpu(hostMean.ravel()), GPUArray.toGpu(hostVar.ravel()) outdata = context.batchNormNd(data, mean, var, scale, bias, test=True) hostOutData = ((hostData - hostMean) / np.sqrt(hostVar + epsilon) * hostScale + hostBias).astype(dtype) assert np.allclose(hostOutData, outdata.get(), atol=atol)
def ctcLoss(data, datalen, labels, lengths, blank, error=None, normalized=False, returnAlphas=False): T, batchsize, vocabsize = data.shape mx = 2 * np.max(lengths) + 1 config = min(i for i, (NT, VT) in enumerate(configs) if mx <= NT * VT) mod, NT = modules[config], configs[config][0] if not normalized: data = cudnn.softmaxNd(data.reshape(T * batchsize, vocabsize, 1, 1), allocator=memPool).reshape( T, batchsize, vocabsize ) offsets = np.cumsum(lengths, dtype=np.int32) extOffsets = np.empty(shape=(batchsize + 1, ), dtype=np.int32) extOffsets[0] = 0 extOffsets[1:] = offsets alphas = GPUArray.empty((T * (2 * int(offsets[-1]) + batchsize), ), dtype=np.float32, allocator=memPool) offsets = GPUArray.toGpu(extOffsets, allocator=memPool) nll = GPUArray.empty((batchsize, ), dtype=np.float32, allocator=memPool) error = GPUArray.zeros((), dtype=np.float32, allocator=memPool) if error is None else error grad = GPUArray.zeros(data.shape, dtype=np.float32, allocator=memPool) mod.calcAlphas( data, datalen, np.int32(T), np.int32(vocabsize), labels, offsets, alphas, np.int32(blank), nll, error, block=(NT, 1, 1), grid=(batchsize, 1, 1) ) mod.calcBetas( data, datalen, np.int32(T), np.int32(vocabsize), labels, offsets, alphas, np.int32(blank), nll, grad, block=(NT, 1, 1), grid=(batchsize, 1, 1) ) return (error, grad) if not returnAlphas else (error, grad, alphas)
def embed(data, W, allocator=memPool):
assert data.dtype == np.int32 and W.dtype == np.float32
batchsize, sentlen = data.shape
_, embsize = W.shape
outdata = GPUArray.zeros((batchsize, sentlen, embsize),
dtype=np.float32,
allocator=allocator)
size = batchsize * sentlen
block = (warpSize, warpSize, 1)
grid = (roundUpDiv(embsize, warpSize), roundUpDiv(size, warpSize), 1)
mod.embed(outdata,
data,
W,
np.int32(size),
np.int32(embsize),
block=block,
grid=grid)
return outdata