最近在使用pytorch编写transformer模型时遇到一个奇怪的报错,使用CPU训练模型时没有问题,但是一切换到GPU训练时就出问题,会报
RuntimeError: CUDA error: an illegal memory access was encountered
- 首先在文件头插入以下代码,可以让报错信息更准确
os.environ['CUDA_LAUNCH_BLOCKING'] = '1'
- 之后再运行模型,发现报错在forward函数的最后一行上
return nn.LayerNorm(d_model)(output + residual)
这行代码上,经过查询资料得知,这是因为nn.LayerNorm这个layer还储存在cpu上,要计算cuda上保存的变量时就会报错。所以正确的解决方法就是在构造模型时就声明nn.LayerNorm,不要在forward中声明nn.layerNorm
class transformer(nn.Module):
def __init__(self):
###代码###
xxxxx
##########
self.layerNorm = nn.LayerNorm(d_model)
def forward(self):
###代码###
xxxxx
##########
return self.layerNorm(output + residual)
希望我的经验能帮助到你 =w=