Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Use cudnn.convert to convert from GPU to CPU #96

Open
wants to merge 1 commit into
base: master
Choose a base branch
from
Open

Use cudnn.convert to convert from GPU to CPU #96

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
49 changes: 1 addition & 48 deletions convert_checkpoint_gpu_to_cpu.lua
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,67 +23,20 @@ cmd:text('Convert a GPU checkpoint to CPU checkpoint.')
cmd:text()
cmd:text('Options')
cmd:argument('-model','GPU model checkpoint to convert')
cmd:option('-gpuid',0,'which gpu to use. -1 = use CPU')
cmd:text()

-- parse input params
local opt = cmd:parse(arg)
torch.manualSeed(123)
torch.setdefaulttensortype('torch.FloatTensor') -- for CPU
cutorch.setDevice(opt.gpuid + 1) -- note +1 because lua is 1-indexed

local checkpoint = torch.load(opt.model)
local protos = checkpoint.protos

-------------------------------------------------------------------------------
-- these functions are adapted from Michael Partheil
-- https://groups.google.com/forum/#!topic/torch7/i8sJYlgQPeA
-- the problem is that you can't call :float() on cudnn module, it won't convert
function replaceModules(net, orig_class_name, replacer)
local nodes, container_nodes = net:findModules(orig_class_name)
for i = 1, #nodes do
for j = 1, #(container_nodes[i].modules) do
if container_nodes[i].modules[j] == nodes[i] then
local orig_mod = container_nodes[i].modules[j]
print('replacing a cudnn module with nn equivalent...')
print(orig_mod)
container_nodes[i].modules[j] = replacer(orig_mod)
end
end
end
end
function cudnnNetToCpu(net)
local net_cpu = net:clone():float()
replaceModules(net_cpu, 'cudnn.SpatialConvolution',
function(orig_mod)
local cpu_mod = nn.SpatialConvolution(orig_mod.nInputPlane, orig_mod.nOutputPlane,
orig_mod.kW, orig_mod.kH, orig_mod.dW, orig_mod.dH, orig_mod.padW, orig_mod.padH)
cpu_mod.weight:copy(orig_mod.weight)
cpu_mod.bias:copy(orig_mod.bias)
cpu_mod.gradWeight = nil -- sanitize for thinner checkpoint
cpu_mod.gradBias = nil -- sanitize for thinner checkpoint
return cpu_mod
end)
replaceModules(net_cpu, 'cudnn.SpatialMaxPooling',
function(orig_mod)
local cpu_mod = nn.SpatialMaxPooling(orig_mod.kW, orig_mod.kH, orig_mod.dW, orig_mod.dH,
orig_mod.padW, orig_mod.padH)
return cpu_mod
end)
replaceModules(net_cpu, 'cudnn.ReLU', function() return nn.ReLU() end)
return net_cpu
end
-------------------------------------------------------------------------------

-- convert the networks to be CPU models
for k,v in pairs(protos) do
print('converting ' .. k .. ' to CPU')
if k == 'cnn' then
-- the cnn is a troublemaker
local cpu_cnn = cudnnNetToCpu(v)
protos[k] = cpu_cnn
protos[k] = cudnn.convert(v, nn):float()
elseif k == 'lm' then
local debugger = require('fb.debugger'); debugger:enter()
v.clones = nil -- sanitize the clones inside the language model (if present just in case. but they shouldnt be)
v.lookup_tables = nil
protos[k]:float() -- ship to CPU
Expand Down