openface/training/model.lua

51 lines
1.1 KiB
Lua
Raw Normal View History

require 'nn'
require 'dpnn'
require 'optim'
2016-01-12 04:36:58 +08:00
if opt.cuda then
require 'cunn'
if opt.cudnn then
require 'cudnn'
cudnn.benchmark = false
cudnn.fastest = true
cudnn.verbose = false
end
end
paths.dofile('torch-TripletEmbedding/TripletEmbedding.lua')
if 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)
2016-01-12 05:44:50 +08:00
assert(imgDim, "Model definition must set global variable 'imgDim'")
assert(imgDim == opt.imgDim, "Model definiton's imgDim must match imgDim option.")
model = createModel()
end
criterion = nn.TripletEmbeddingCriterion(opt.alpha)
2016-01-12 04:36:58 +08:00
if opt.cuda then
model = model:cuda()
if opt.cudnn then
cudnn.convert(model,cudnn)
end
2016-01-12 04:36:58 +08:00
criterion:cuda()
end
2016-06-17 08:12:16 +08:00
optimizeNet(model, opt.imgDim)
print('=> Model')
print(model)
print(('Number of Parameters: %d'):format(model:getParameters():size(1)))
print('=> Criterion')
print(criterion)
collectgarbage()