Permalink
Switch branches/tags
Nothing to show
Find file Copy path
Fetching contributors…
Cannot retrieve contributors at this time
173 lines (143 sloc) 5.14 KB
require 'torch'
require 'nn'
require 'optim'
require 'image'
require 'fast_neural_style.PerceptualCriterion'
require 'fast_neural_style.TotalVariation'
local utils = require 'fast_neural_style.utils'
local preprocess = require 'fast_neural_style.preprocess'
--[[
Perform optimization-based style transfer as described in
--]]
local cmd = torch.CmdLine()
-- Basic options
cmd:option('-content_image', 'images/content/chicago.jpg')
cmd:option('-style_image', 'images/styles/starry_night.jpg')
cmd:option('-image_size', 512)
-- Loss options
cmd:option('-loss_network', 'models/vgg16.t7')
cmd:option('-tv_strength', 1e-6)
cmd:option('-loss_type', 'L2', 'L2|SmoothL1')
cmd:option('-style_target_type', 'gram', 'gram|mean')
-- Options for content reconstruction
cmd:option('-content_weights', '1.0')
cmd:option('-content_layers', '16')
-- Options for style reconstruction
cmd:option('-style_weights', '5.0')
cmd:option('-style_layers', '4,9,16,23')
cmd:option('-style_image_size', 512)
-- Options for DeepDream
cmd:option('-deepdream_layers', '')
cmd:option('-deepdream_weights', '')
-- Optimization
cmd:option('-learning_rate', 1.0)
cmd:option('-optimizer', 'lbfgs', 'lbfgs|adam')
cmd:option('-num_iterations', 500)
-- Output options
cmd:option('-output_image', 'out.png')
cmd:option('-print_every', 1)
cmd:option('-save_every', 50)
-- Other options
cmd:option('-preprocessing', 'vgg')
-- Backend options
cmd:option('-gpu', -1)
cmd:option('-backend', 'cuda', 'cuda|opencl')
cmd:option('-use_cudnn', 1)
local opt = cmd:parse(arg)
local function main()
if not preprocess[opt.preprocessing] then
local msg = 'invalid -preprocessing "%s"; must be "vgg" or "resnet"'
error(string.format(msg, opt.preprocessing))
end
preprocess = preprocess[opt.preprocessing]
local dtype, use_cudnn = utils.setup_gpu(opt.gpu, opt.backend, opt.use_cudnn == 1)
-- Set up the criterion
local ok, loss_net = pcall(function() return torch.load(opt.loss_network) end)
if not ok then
print('ERROR: Could not load loss network from ' .. opt.loss_network)
print('You may need to download the VGG-16 model by running:')
print('bash models/download_vgg16.sh')
return
end
print(loss_net)
local style_layers, style_weights =
utils.parse_layers(opt.style_layers, opt.style_weights)
local content_layers, content_weights =
utils.parse_layers(opt.content_layers, opt.content_weights)
local deepdream_layers, deepdream_weights =
utils.parse_layers(opt.deepdream_layers, opt.deepdream_weights)
local crit_args = {
cnn = loss_net,
style_layers = style_layers,
style_weights = style_weights,
content_layers = content_layers,
content_weights = content_weights,
deepdream_layers = deepdream_layers,
deepdream_weights = deepdream_weights,
loss_type = opt.loss_type,
agg_type = opt.style_target_type,
}
local crit = nn.PerceptualCriterion(crit_args):type(dtype)
-- Set the content image
local content_image = image.load(opt.content_image, 3)
content_image = image.scale(content_image, opt.image_size)
local H, W = content_image:size(2), content_image:size(3)
content_image = preprocess.preprocess(content_image:view(1, 3, H, W))
crit:setContentTarget(content_image:type(dtype))
-- Set the style image
local style_image = image.load(opt.style_image, 3)
style_image = image.scale(style_image, opt.style_image_size)
local H, W = style_image:size(2), style_image:size(3)
style_image = preprocess.preprocess(style_image:view(1, 3, H, W))
crit:setStyleTarget(style_image:type(dtype))
-- Set up total variation regularization
local tv = nn.Identity()
if opt.tv_strength > 0 then
tv = nn.TotalVariation(opt.tv_strength)
end
tv:type(dtype)
local img = torch.randn(#content_image):type(dtype)
-- Callback function for optim methods
local f_calls = 0
local function f(x)
f_calls = f_calls + 1
local tv_out = tv:forward(x)
local loss = crit:forward(tv_out, {})
local grad_tv_out = crit:updateGradInput(tv_out, {})
local grad_x = tv:backward(x, grad_tv_out)
if opt.print_every > 0 and f_calls % opt.print_every == 0 then
print(string.format('Iteration %d, loss = %f', f_calls, loss))
end
if opt.save_every > 0 and f_calls % opt.save_every == 0 then
local img_out = preprocess.deprocess(img:float())[1]
local ext = paths.extname(opt.output_image)
local basename = paths.basename(opt.output_image):split('%.')[1]
local directory = paths.dirname(opt.output_image)
local filename = string.format('%s/%s_%d.%s',
directory, basename, f_calls, ext)
image.save(filename, img_out)
end
return loss, grad_x:view(-1)
end
if opt.optimizer == 'lbfgs' then
local config = {
maxIter=opt.num_iterations,
learningRate=opt.learning_rate,
verbose=true,
tolX = 0,
}
optim.lbfgs(f, img, config)
else
local config = {
learningRate = opt.learning_rate
}
local optim_fn = optim[opt.optimizer]
if optim_fn == nil then
error(string.format('Invalid optimizer "%s"', opt.optimizer))
end
for t = 1, opt.num_iterations do
optim_fn(f, img, config)
end
end
end
main()