openface/training/model.lua

68 lines
1.4 KiB
Lua

require 'nn'
require 'dpnn'
require 'optim'
if opt.cuda then
require 'cunn'
if opt.cudnn then
require 'cudnn'
cudnn.benchmark = opt.cudnn_bench
cudnn.fastest = true
cudnn.verbose = false
end
end
paths.dofile('torch-TripletEmbedding/TripletEmbedding.lua')
local M = {}
function M.modelSetup(continue)
if continue then
model = continue
elseif opt.retrain ~= 'none' then
assert(paths.filep(opt.retrain), 'File not found: ' .. opt.retrain)
print('Loading model from file: ' .. opt.retrain);
model = torch.load(opt.retrain)
print("Using imgDim = ", opt.imgDim)
else
paths.dofile(opt.modelDef)
assert(imgDim, "Model definition must set global variable 'imgDim'")
assert(imgDim == opt.imgDim, "Model definiton's imgDim must match imgDim option.")
model = createModel()
end
-- First remove any DataParallelTable
if torch.type(model) == 'nn.DataParallelTable' then
model = model:get(1)
end
criterion = nn.TripletEmbeddingCriterion(opt.alpha)
if opt.cuda then
model = model:cuda()
if opt.cudnn then
cudnn.convert(model,cudnn)
end
criterion:cuda()
else
model:float()
criterion:float()
end
optimizeNet(model, opt.imgDim)
if opt.cuda and opt.nGPU > 1 then
model = makeDataParallel(model, opt.nGPU)
end
collectgarbage()
return model, criterion
end
return M