Permalink
Find file
73e1715 Nov 25, 2015
99 lines (90 sloc) 3.6 KB
--[[
A quick script for converting GPU checkpoints to CPU checkpoints.
CPU checkpoints are not saved by the training script automatically
because of Torch cloning limitations. In particular, it is not
possible to clone a GPU model on CPU, something like :clone():float()
with a single call, without needing extra memory on the GPU. If this
existed then it would be possible to do this inside the training
script without worrying about blowing up the memory.
]]--
require 'torch'
require 'nn'
require 'nngraph'
require 'cutorch'
require 'cunn'
require 'cudnn' -- only needed if the loaded model used cudnn as backend. otherwise can be commented out
-- local imports
require 'misc.LanguageModel'
cmd = torch.CmdLine()
cmd:text()
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
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
else
error('error: strange module in protos: ' .. k)
end
end
local savefile = opt.model .. '_cpu.t7' -- append "cpu.t7" to filename
torch.save(savefile, checkpoint)
print('saved ' .. savefile)