-
Notifications
You must be signed in to change notification settings - Fork 69
/
convert_gpu_to_cpu.lua
45 lines (36 loc) · 1.04 KB
/
convert_gpu_to_cpu.lua
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
require 'torch';
require 'nn';
cmd = torch.CmdLine()
cmd:option('-loadPath', 'checkpoints/model.t7')
cmd:option('-savePath', 'checkpoints/model_cpu.t7')
cmd:option('-gpuid', 0)
opt = cmd:parse(arg)
-- check for new save path
if opt.savePath == 'checkpoints/model_cpu.t7' then
opt.savePath = opt.loadPath .. '.cpu.t7'
end
print(opt)
if opt.gpuid >= 0 then
require 'cutorch'
require 'cunn'
if opt.backend == 'cudnn' then require 'cudnn' end
cutorch.setDevice(opt.gpuid+1)
torch.setdefaulttensortype('torch.CudaTensor');
else
print('Gotta have a GPU to convert to CPU :(')
os.exit()
end
print('Loading model')
model = torch.load(opt.loadPath)
-- convert modelW and optims to cpu
print('Shipping params to CPU')
if model.modelW:type() == 'torch.CudaTensor' then
model.modelW = model.modelW:float()
end
for k,v in pairs(model.optims) do
if torch.type(v) ~= 'number' and v:type() == 'torch.CudaTensor' then
model.optims[k] = v:float()
end
end
print('Saving to ' .. opt.savePath)
torch.save(opt.savePath, model)