Allow optimization on CPU, adjust learning rate for VGG so that SGD doesn't diverge

Author: Kai Sheng Tai

.gitignore

@@ -1,3 +1,4 @@
 .DS_Store
 frames
-*~
+*~
+*.th

costs.lua

@@ -36,7 +36,11 @@ end
 
 local euclidean = nn.MSECriterion()
 euclidean.sizeAverage = false
-euclidean:cuda()
+if opt.cpu then
+    euclidean:float()
+else
+    euclidean:cuda()
+end
 
 function style_grad(gen, orig_gram)
     local k = gen:size(2)

images.lua

@@ -23,7 +23,8 @@ function preprocess(img, scale)
         end
     end
 
-    local copy = torch.Tensor(img:size())
+    -- reverse channels
+    local copy = torch.FloatTensor(img:size())
     copy[1] = img[3]
     copy[2] = img[2]
     copy[3] = img[1]
@@ -37,17 +38,17 @@ function preprocess(img, scale)
 end
 
 function depreprocess(img)
-    img = img:float():view(3, img:size(3), img:size(4))
+    local copy = torch.FloatTensor(3, img:size(3), img:size(4)):copy(img)
     for i = 1, 3 do
-        img[i]:add(means[i])
+        copy[i]:add(means[i])
     end
-    img:div(255)
+    copy:div(255)
 
-    local copy = torch.FloatTensor(img:size())
-    copy[1] = img[3]
-    copy[2] = img[2]
-    copy[3] = img[1]
-    img = copy
-    img:clamp(0, 1)
-    return img
+    -- reverse channels
+    local copy2 = torch.FloatTensor(copy:size())
+    copy2[1] = copy[3]
+    copy2[2] = copy[2]
+    copy2[3] = copy[1]
+    copy2:clamp(0, 1)
+    return copy2
 end

main.lua

@@ -31,6 +31,7 @@ cmd:option('--smoothness',       0,       'Total variation norm regularization s
 cmd:option('--init',            'image',  '{image, random}. Initialization mode for optimized image.')
 cmd:option('--backend',         'cunn',   '{cunn, cudnn}. Neural network CUDA backend.')
 cmd:option('--optimizer',       'lbfgs',  '{sgd, lbfgs}. Optimization algorithm.')
+cmd:option('--cpu',              false,   'Optimize on CPU (only with VGG network).')
 opt = cmd:parse(arg)
 if opt.size <= 0 then
     opt.size = nil
@@ -51,6 +52,10 @@ if opt.model == 'inception' then
         print('run download_models.sh to download model weights')
         error('')
     end
+
+    if opt.cpu then
+        error('CPU optimization only works with VGG model')
+    end
 elseif opt.model == 'vgg' then
     if not paths.filep(vgg_path) then
         print('ERROR: could not find VGG model weights at ' .. vgg_path)
@@ -97,6 +102,13 @@ elseif opt.model == 'vgg' then
 
     model = create_vgg(vgg_path, opt.backend)
 end
+
+-- run on GPU
+if opt.cpu then
+    model:float()
+else
+    model:cuda()
+end
 collectgarbage()
 
 -- compute normalization factor
@@ -111,14 +123,20 @@ for k, v in pairs(content_weights) do
 end
 
 -- load content image
-local img = preprocess(image.load(opt.content), opt.size):cuda()
+local img = preprocess(image.load(opt.content), opt.size)
+if not opt.cpu then
+    img = img:cuda()
+end
 model:forward(img)
 local img_activations, _ = collect_activations(model, content_weights, {})
 
 -- load style image
 local art = preprocess(
     image.load(opt.style), math.max(img:size(3), img:size(4))
-):cuda()
+)
+if not opt.cpu then
+    art = art:cuda()
+end
 model:forward(art)
 local _, art_grams = collect_activations(model, {}, style_weights)
 art = nil
@@ -130,7 +148,8 @@ function opfunc(input)
 
     -- backpropagate
     local loss = 0
-    local grad = torch.CudaTensor(model.output:size()):zero()
+    local grad = opt.cpu and torch.FloatTensor() or torch.CudaTensor()
+    grad:resize(model.output:size()):zero()
     for i = #model.modules, 1, -1 do
         local module_input = (i == 1) and input or model.modules[i - 1].output
         local module = model.modules[i]
@@ -168,7 +187,11 @@ if opt.init == 'image' then
 elseif opt.init == 'random' then
     input = preprocess(
         torch.randn(3, img:size(3), img:size(4)):mul(0.1):add(0.5):clamp(0, 1)
-    ):cuda()
+    )
+
+    if not opt.cpu then
+        input = input:cuda()
+    end
 else
     error('unrecognized initialization option: ' .. opt.init)
 end
@@ -190,10 +213,15 @@ image.save(paths.concat(frames_dir, '0.jpg'), output)
 local optim_state
 if opt.optimizer == 'sgd' then
     optim_state = {
-        learningRate = 0.1,
         momentum = 0.9,
         dampening = 0.0,
     }
+
+    if opt.model == 'inception' then
+        optim_state.learningRate = 5e-2
+    else
+        optim_state.learningRate = 1e-3
+    end
 elseif opt.optimizer == 'lbfgs' then
     optim_state = {
         maxIter = 3,

models/inception.lua

@@ -97,6 +97,5 @@ function create_inception(weights_file, backend)
 
     model = model:subnetwork('inception_4e')
     collectgarbage()
-    model:cuda()
     return model
 end

models/vgg19.lua

@@ -63,6 +63,5 @@ function create_vgg(weights_file, backend)
     end
 
     collectgarbage()
-    model:cuda()
     return model
 end