Tensor Comprehensions 基础算子改写

最新推荐文章于 2024-08-06 04:53:13 发布

原创最新推荐文章于 2024-08-06 04:53:13 发布 · 570 阅读

0 ·

CC 4.0 BY-SA版权

云计算专栏收录该内容

3 篇文章

订阅专栏

本文通过定义自定义卷积、ReLU激活、池化、批量归一化及全连接层等核心模块，详细介绍了如何使用TorchCraft库实现深度学习网络的基本组件。这些组件包括卷积层、ReLU激活函数、最大池化、平均池化、批量归一化层及线性层。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

def conv(self, input, outchannel, k_size, stride = 1, padding = 0):
LANG = """
def convolution(float(N,C,H,W) I, float(M,C,KH,KW) W1) -> (O) {{
O(n, m, h, w) +=! I(n, c, {sh} * h + kh, {sw} * w + kw) * W1(m, c, kh, kw)
}}
"""
kernel = torch.randn(outchannel, input.size(1), k_size, k_size)
input = input.type(torch.FloatTensor)
if padding != 0:
pad = torch.zeros(input.size(0), input.size(1), padding, input.size(3))
input = torch.cat((input, pad), 2)
input = torch.cat((pad, input), 2)
pad = torch.zeros(input.size(0), input.size(1), input.size(2), padding)
input = torch.cat((input, pad), 3)
input = torch.cat((pad, input), 3)

sh, sw = stride, stride
convolution = tc.define(LANG, training=True, name="convolution", backward="convolution_grad",constants={"sh":sh, "sw":sw})
I = Variable(input.cuda(), requires_grad=True)
W = Parameter(kernel.cuda())
out = convolution(I, W, options=tc.CudaMappingOptions("conv"))
return out

def relu(self, input):
LANG = """
def relu(float(Q, W, B, M) I) -> (O1){
O1(q, w, b, m) = fmax(I(q, w, b, m), 0)
}
"""
relu = tc.define(LANG, name="relu")
inp = input.cuda()
out = relu(inp, options=tc.CudaMappingOptions("naive"))
return out

def maxpool(self, input, p_size, stride, padding = 0):
LANG = """
def maxpool(float(B,C,H,W) input) -> (output) {{
output(b,c,h,w) max=! input(b, c, h * {sH} + kh, w * {sW} + kw)
where kh in 0:{kH}, kw in 0:{kW}
}}
"""
input = input.type(torch.FloatTensor)
if padding != 0:
pad = torch.zeros(input.size(0), input.size(1), padding, input.size(3))
input = torch.cat((input, pad), 2)
input = torch.cat((pad, input), 2)
pad = torch.zeros(input.size(0), input.size(1), input.size(2), padding)
input = torch.cat((input, pad), 3)
input = torch.cat((pad, input), 3)

sH, sW = stride, stride
kH, kW = p_size, p_size
maxpool = tc.define(LANG, name="maxpool", constants={"sH":sH, "sW":sW, "kH":kH, "kW":kW})
inp = input.cuda()
out = maxpool(inp, options=tc.CudaMappingOptions("naive"))
out = out.type(torch.FloatTensor)
return out

def avgpool(self, input, p_size, stride, padding = 0):
LANG = """
def avgpool(float(B, C, H, W) input) -> (output) {{
output(b, c, h, w) +=! input(b, c, h * {sH} + kh, w * {sW} + kw) / ({kH} * {kW})
where kh in 0:{kH}, kw in 0:{kW}
}}
"""
input = input.type(torch.FloatTensor)
if padding != 0:
pad = torch.zeros(input.size(0), input.size(1), padding, input.size(3))
input = torch.cat((input, pad), 2)
input = torch.cat((pad, input), 2)
pad = torch.zeros(input.size(0), input.size(1), input.size(2), padding)
input = torch.cat((input, pad), 3)
input = torch.cat((pad, input), 3)

sH, sW = stride, stride
kH, kW = p_size, p_size
avgpool = tc.define(LANG, name="avgpool", constants={"sH":sH, "sW":sW, "kH":kH, "kW":kW})
inp = input.cuda()
out = avgpool(inp, options=tc.CudaMappingOptions("naive"))
out = out.type(torch.FloatTensor)
return out

def batchnorm(self, input):
LANG = """
def batchnorm(float(N,C,H,W) I, float(C) rMeanIn, float(C) rVarIn)
-> (O, rMeanOut, rVarOut, mean, centered, variance, expectedVariance, normalizedOut)
{{
mean(c) +=! I(nn, c, hh, ww)
mean(c) = mean(c) / (N * H * W)
rMeanOut(c) = (1 - {momentum}) * rMeanIn(c) + {momentum} * mean(c)
centered(n, c, h, w) = I(n, c, h, w) - rMeanOut(c)
variance(n, c, h, w) = centered(n, c, h, w) * centered(n, c, h, w)
expectedVariance(c) +=! (variance(n, c, h, w) + {eps}) / (N * H * W)
rVarOut(c) = rsqrt(
(1 - {momentum}) * rVarIn(c) + {momentum} * expectedVariance(c))
O(n, c, h, w) = centered(n, c, h, w) * rVarOut(c)
normalizedOut(n, c, h, w) = O(n, c, h, w)
}}
"""
channel = input.size(1)
batchnorm = tc.define(LANG, name="batchnorm", constants={"momentum": 0.9, "eps": 1e-5})
running_mean, running_var = torch.randn(channel).cuda(), torch.randn(channel).cuda()
out = batchnorm(input.cuda(), running_mean, running_var, options=tc.CudaMappingOptions("naive"))
return out[7]

def linear(self, input, outchannel):
LANG = """
def fc(float(B,M) I, float(N,M) W1, float(N) B1) -> (O1) {
O1(b, n) +=! I(b, m) * W1(n, m)
O1(b, n) = O1(b, n) + B1(n)
}
"""
I = input.cuda()
W = torch.randn(outchannel, input.size(1)).cuda()
B = torch.randn(outchannel).cuda()
fc = tc.define(LANG, name="fc")
out = fc(I, W, B, options=tc.CudaMappingOptions("naive"))
return out