【RuntimeError: CUDA error: out of memory】pytorch4.0+py3.6+[LinearStyleTransfer项目]

记录项目LinearStyleTransfer所遇到的问题:

附上项目链接:LinearStyleTransfer

按照github所述,安装好环境(我的环境)

pytorch0.4.1

python3.6

cuda8.0

opencv3.4

所需模型:models

一、环境准备妥当,运行测试文件TestArtistic.py时,报错:

【RuntimeError: CUDA error: out of memory】pytorch4.0+py3.6+[LinearStyleTransfer项目]

查了一些原因,有说是pytorch版本高才出的问题,但是我的pytorch版本与作者的是一样的,所以觉得不是这个原因

然后在代码中加了一些log,发现内存超限的位置是:加载模型 时
【RuntimeError: CUDA error: out of memory】pytorch4.0+py3.6+[LinearStyleTransfer项目]

有个博主分析说是,有些项目导出时会将显卡信息一并导出,而本来默认的gpu_id为0,但是我们的gpu上0号显卡被占用,所以要设置下模型的加载位置,改动方法:

(1)加入gpu_id,默认值随意

         parser.add_argument("--gpu_id", type=int, default=1,help='which gpu to use')

【RuntimeError: CUDA error: out of memory】pytorch4.0+py3.6+[LinearStyleTransfer项目]

(2)if(opt.cuda):

              torch.cuda.set_device(opt.gpu_id)

【RuntimeError: CUDA error: out of memory】pytorch4.0+py3.6+[LinearStyleTransfer项目]

(3)在加载模型时设置下加载位置:(我这里是1号显卡)

matrix.load_state_dict(torch.load(opt.matrixPath, map_location=lambda storage, loc: storage.cuda(1)))

【RuntimeError: CUDA error: out of memory】pytorch4.0+py3.6+[LinearStyleTransfer项目]

再次运行测试文件可通过。。。