Repository: 1adrianb/face-alignment-training Branch: master Commit: 1e3127d9a407 Files: 13 Total size: 32.8 KB Directory structure: gitextract_3ng_0kzy/ ├── .gitignore ├── LICENSE ├── README.md ├── checkpoints.lua ├── dataloader.lua ├── dataset-images.lua ├── dataset-init.lua ├── main.lua ├── models/ │ ├── fan.lua │ └── init.lua ├── opts.lua ├── train.lua └── utils.lua ================================================ FILE CONTENTS ================================================ ================================================ FILE: .gitignore ================================================ # Compiled Lua sources luac.out # luarocks build files *.src.rock *.zip *.tar.gz # Object files *.o *.os *.ko *.obj *.elf # Precompiled Headers *.gch *.pch # Libraries *.lib *.a *.la *.lo *.def *.exp # Shared objects (inc. Windows DLLs) *.dll *.so *.so.* *.dylib # Executables *.exe *.out *.app *.i*86 *.x86_64 *.hex ================================================ FILE: LICENSE ================================================ This pipeline is build around the ImageNet training code avaialable at and HourGlass(HG) code available at https://github.com/anewell/pose-hg-train Copyright (c) 2016, Facebook, Inc. Copyright (c) 2016, University of Michigan For the rest of the code and models: Copyright (c) 2017, University of Nottingham All rights reserved. Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met: 1.Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer. 2.Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution. 3.Neither the name of the paper nor the names of its contributors may be used to endorse or promote products derived from this software without specific prior written permission. THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. ================================================ FILE: README.md ================================================ # How far are we from solving the 2D \& 3D Face Alignment problem? (and a dataset of 230,000 3D facial landmarks) This is the training code for 2D-FAN and 3D-FAN decribed in "How far are we from solving the 2D \& 3D Face Alignment problem? (and a dataset of 230,000 3D facial landmarks)" paper. Please visit [our](https://www.adrianbulat.com) webpage or read bellow for instructions on how to run the code. Pretrained models are available on our page. **Demo code: ** Note: If you are interested in a binarized version, capable of running on devices with limited resources please also check for a demo. ## Requirments - Install the latest [Torch7](http://torch.ch/docs/getting-started.html) version (for Windows, please follow the instructions available [here](https://github.com/torch/distro/blob/master/win-files/README.md)) ### Packages - [cutorch](https://github.com/torch/cutorch) - [nn](https://github.com/torch/nn) - [nngraph](https://github.com/torch/nngraph) - [cudnn](https://github.com/soumith/cudnn.torch) - [xlua](https://github.com/torch/xlua) - [image](https://github.com/torch/image) - [paths](https://github.com/torch/paths) - [matio](https://github.com/soumith/matio-ffi.torch) ## Setup 1. Clone the github repository and install all the dependencies mentiones above. ```bash git clone https://github.com/1adrianb/face-alignment-training cd face-alignment-training ``` 2. Download the 300W-LP dataset from the authors webpage. In order to train on your own data the dataloader.lua file needs to be adapted. 3. Download the 300W-LP annotations converted to t7 format from [here](https://www.adrianbulat.com/downloads/FaceAlignment/landmarks.zip), extract it and move the ```landmarks``` folder to the root of the 300W-LP dataset. ## Usage In order to run the demo please download the required models available bellow and the associated data. ```bash th main.lua -data path_to_300W_LP_dataset ``` In order to see all the available options please run: ```bash th main.lua --help ``` ## Citation ``` @inproceedings{bulat2017far, title={How far are we from solving the 2D \& 3D Face Alignment problem? (and a dataset of 230,000 3D facial landmarks)}, author={Bulat, Adrian and Tzimiropoulos, Georgios}, booktitle={International Conference on Computer Vision}, year={2017} } ``` ## Acknowledgements This pipeline is build around the ImageNet training code avaialable at and HourGlass(HG) code available at https://github.com/anewell/pose-hg-train ================================================ FILE: checkpoints.lua ================================================ local checkpoint = {} function checkpoint.latest(opt) if opt.resume == 'none' then return nil end local latestPath = paths.concat(opt.resume, 'latest.t7') if not paths.filep(latestPath) then return nil end print('=> Loading checkpoint ' .. latestPath) local latest = torch.load(latestPath) local optimState = torch.load(paths.concat(opt.resume, latest.optimFile)) return latest, optimState end function checkpoint.save(epoch, model, optimState) -- Don't save the DataParallelTable for easier loading on other machines if torch.type(model) == 'nn.DataParallelTable' then model = model:get(1) end local modelFile = 'model_' .. epoch .. '.t7' local optimFile = 'optimState_' .. epoch .. '.t7' torch.save(modelFile, model) torch.save(optimFile, optimState) torch.save('latest.t7', { epoch = epoch, modelFile = modelFile, optimFile = optimFile, }) end return checkpoint ================================================ FILE: dataloader.lua ================================================ local datasets = require 'dataset-init' local Threads = require 'threads' Threads.serialization('threads.sharedserialize') local M = {} local DataLoader = torch.class('DataLoader', M) function DataLoader.getDataFaces(opt, split) print('=> Building dataset...') base_dir = opt.data..'landmarks/' dirs = paths.dir(base_dir) lines = {} for i=1,#dirs do if string.sub(dirs[i],1,1) ~= '.' then for f in paths.files(base_dir..dirs[i],'.mat') do if not string.find(f, "test") then lines[#lines+1] = f end end end end print('=> Dataset built. '..#lines..' images were found.') return lines end function DataLoader.create(opt,split) local _dataset = nil local dataAnnot = DataLoader.getDataFaces(opt,split) return M.DataLoader(_dataset,opt,split,dataAnnot) end function DataLoader:__init(_dataset, opt, split, dataAnnot) local manualSeed = opt.manualSeed local function init() local datasets = require 'dataset-init' trainLoader, valLoader = datasets.create(opt,split,dataAnnot) end local function main(idx) if manualSeed ~= 0 then torch.manualSeed(manualSeed + idx) end torch.setnumthreads(1) _G.dataset = trainLoader return trainLoader:size() end local threads, sizes = Threads(opt.nThreads,init, main) self.nCrops = 1 self.threads = threads self.__size = sizes[1][1] self.batchSize = math.floor(opt.batchSize / self.nCrops) self.opt = opt self.split = split self.dataAnnot = dataAnnot end function DataLoader:size() return math.ceil(self.__size/self.batchSize) end function DataLoader:annot() trainLoader = datasets.create(self.opt,self.split,self.dataAnnot) return trainLoader.annot end function DataLoader:run() local threads = self.threads local size, batchSize = self.__size, self.batchSize local perm = torch.randperm(size) local idx, sample = 1, nil local function enqueue() while idx <= size and threads:acceptsjob() do local indices = perm:narrow(1,idx,math.min(batchSize,size-idx+1)) threads:addjob( function(indices,nCrops) local sz = indices:size(1) local batch, imageSize local target local indicesCopy = indices for i,idx in ipairs(indices:totable()) do local sample, label = _G.dataset:get(false,idx) local input, label = _G.dataset:preprocess(sample, label) if not batch then imageSize = input:size():totable() if nCrops > 1 then table.remove(imageSize,1) end batch = torch.FloatTensor(sz,nCrops, table.unpack(imageSize)) end if not target then targetSize = label:size():totable() target = torch.FloatTensor(sz,nCrops, table.unpack(targetSize)) end batch[i]:copy(input) target[i]:copy(label) end collectgarbage() return { input = batch:view(sz*nCrops,table.unpack(imageSize)), label = target:view(sz*nCrops,table.unpack(targetSize)), indx = indicesCopy , } end, function(_sample_) sample = _sample_ end, indices, self.nCrops ) idx = idx + batchSize end end local n = 0 local function loop() enqueue() if not threads:hasjob() then return nil end threads:dojob() if threads:haserror() then threads:synchronize() end enqueue() n = n+1 return n, sample end return loop end return M.DataLoader ================================================ FILE: dataset-images.lua ================================================ local image = require('image') require 'utils' local M = {} local DatasetImages = torch.class('DatasetImages', M) function DatasetImages:__init( opt, split, annot ) self.total = #annot self.nParts = 68 self.annot = annot self.opt = opt self.typeOfData = split end function DatasetImages:generateSampleFace(idx) local main_pts = torch.load(self.opt.data..'landmarks/'..self.annot[idx]:split('_')[1]..'/'..string.sub(self.annot[idx],1,#self.annot[idx]-4)..'.t7') local pts = main_pts[1] --- 2:3D local c = torch.Tensor{450/2,450/2+50} local s = 1.8 local img = image.load(self.opt.data..self.annot[idx]:split('_')[1]..'/'..string.sub(self.annot[idx],1,#self.annot[idx]-8)..'.jpg') local inp = crop(img, c, s, 0, 256) local out = torch.zeros(self.nParts, 64, 64) for i = 1, self.nParts do if pts[i][1] > 0 then -- Checks that there is a ground truth annotation drawGaussian(out[i], transform(torch.add(pts[i],1), c, s, 0, 64), 1) end end return inp,out,pts,c,s end function DatasetImages:get(shuffle,i) local inp, out, pts, c, s = self:generateSampleFace(i) self.pts, self.c, self.s = pts,c,s return inp, out end function DatasetImages:size() return self.total end function DatasetImages:preprocess(input, label) if self.typeOfData == 'train' then local s = torch.randn(1):mul(self.opt.scaleFactor):add(1):clamp(1-self.opt.scaleFactor,1+self.opt.scaleFactor)[1] local r = torch.randn(1):mul(self.opt.rotFactor):clamp(-2*self.opt.rotFactor,2*self.opt.rotFactor)[1] -- Scale/rotation if torch.uniform() <= .6 then r = 0 end local inp,out = 256, 64 local divideBy = 200 input = crop(input, {(inp+1)/2,(inp+1)/2}, inp*s/divideBy, r, inp) label = crop(label, {(out+1)/2,(out+1)/2}, out*s/divideBy, r, out) -- Emulate row resolution if torch.uniform()<=.2 and false then --.35 input = image.scale(input,96,96) input = image.scale(input,256,256) end -- Add jpeg artefacts --[[ if torch.uniform()<=.2 and false then local onlyImg = input[{{1,3},{},{}}] onlyImg = image.compressJPG(onlyImg,30) onlyImg = image.decompressJPG(onlyImg) input[{{1,3},{},{}}] = onlyImg end ]]-- -- Add random translation --[[ wh_t = torch.Tensor(2):random(0,80)-40 input = image.translate(input,wh_t[1],wh_t[2]) label = image.translate(label,wh_t[1]/4.0,wh_t[2]/4.0) ]]-- -- Add some gaussian blue --[[ if torch.uniform()<.4 and false then gauss_s = torch.Tensor(1):random(10,30):int() local kernel_gauss = image.gaussian(gauss_s[1]) input = image.convolve(input, kernel_gauss, 'same')/255.0 end ]]-- local flip_ = customFlip or flip local shuffleLR_ = customShuffleLR or shuffleLR if torch.uniform() <= .5 then input = flip_(input) label = flip_(shuffleLR_(label)) end -- Color augumentation input[{1, {}, {}}]:mul(torch.uniform(0.7, 1.3)):clamp(0, 1) input[{2, {}, {}}]:mul(torch.uniform(0.7, 1.3)):clamp(0, 1) input[{3, {}, {}}]:mul(torch.uniform(0.7, 1.3)):clamp(0, 1) end return input, label end return M.DatasetImages ================================================ FILE: dataset-init.lua ================================================ local M = {} function M.create(opt, split, annot) local Dataset = require('dataset-images') return Dataset(opt, split, annot) end return M ================================================ FILE: main.lua ================================================ require 'torch' require 'cutorch' require 'paths' require 'nn' require 'nngraph' local DataLoader = require 'dataloader' local checkpoints = require 'checkpoints' local models = require 'models/init' local Trainer = require 'train' local opts = require 'opts' local opt = opts.parse(arg) torch.setdefaulttensortype('torch.FloatTensor') torch.setnumthreads(1) cutorch.setDevice(1) torch.setheaptracking(true) torch.manualSeed(opt.manualSeed) cutorch.manualSeed(opt.manualSeed) --Load previous checkpoints, if it exists local checkpoint, optimState = checkpoints.latest(opt) local optimState = checkpoint and torch.load(checkpoint.optimFile) or nil --Create model local model, criterion = models.setup(opt, checkpoint, true) print('=> Model size: ', model:getParameters():size(1)) --Data loading local trainLoader = DataLoader.create(opt,'train') local trainer = Trainer(model, criterion, opt, optimState, netLogger) local startEpoch = checkpoint and checkpoint.epoch + 1 or opt.epochNumber for epoch = startEpoch, opt.nEpochs do -- Train for a single epoch local trainLoss, trainAcc = trainer:train(epoch, trainLoader) print(string.format(' *Results loss: %6.6f acc: %6.6f ',trainLoss, trainAcc)) if opt.snapshot ~= 0 and epoch % opt.snapshot == 0 then checkpoints.save(epoch, model:clearState(), trainer.optimState, bestModel) end end ================================================ FILE: models/fan.lua ================================================ -- Face Alignment Network -- -- How far are we from solving the 2D \& 3D Face Alignment problem? (and a dataset of 230,000 3D facial landmarks) -- Adrian Bulat and Georgios Tzimiropoulos -- ICCV 2017 -- local cudnn = require 'cudnn' -- Define some short names local conv = cudnn.SpatialConvolution local batchnorm = nn.SpatialBatchNormalization local relu = cudnn.ReLU local upsample = nn.SpatialUpSamplingNearest -- Opts local nModules = 1 local nFeats = 256 local nStack = 8 local function convBlock(numIn, numOut, order) local cnet = nn.Sequential() :add(batchnorm(numIn,1e-5,false)) :add(relu(true)) :add(conv(numIn,numOut/2,3,3,1,1,1,1):noBias()) :add(nn.ConcatTable() :add(nn.Identity()) :add(nn.Sequential() :add(nn.Sequential() :add(batchnorm(numOut/2,1e-5,false)) :add(relu(true)) :add(conv(numOut/2,numOut/4,3,3,1,1,1,1):noBias()) ) :add(nn.ConcatTable() :add(nn.Identity()) :add(nn.Sequential() :add(batchnorm(numOut/4,1e-5,false)) :add(relu(true)) :add(conv(numOut/4,numOut/4,3,3,1,1,1,1):noBias()) ) ) :add(nn.JoinTable(2)) ) ) :add(nn.JoinTable(2)) return cnet end -- Skip layer local function skipLayer(numIn,numOut) if numIn == numOut then return nn.Identity() else return nn.Sequential() :add(batchnorm(numIn,1e-5,false)) :add(relu(true)) :add(conv(numIn,numOut,1,1):noBias()) end end -- Residual block local function Residual(numIn,numOut) return nn.Sequential() :add(nn.ConcatTable() :add(convBlock(numIn,numOut)) :add(skipLayer(numIn,numOut))) :add(nn.CAddTable(true)) end local function lin(numIn,numOut,inp) -- Apply 1x1 convolution, stride 1, no padding local l = conv(numIn,numOut,1,1,1,1,0,0)(inp) return relu(true)(batchnorm(numOut)(l)) end local function hourglass(n, f) local model = nn.Sequential() local branch = nn.ConcatTable() local b1 = nn.Sequential() local b2 = nn.Sequential() for i = 1,nModules do b1:add(Residual(f,f)) end b2:add(nn.SpatialMaxPooling(2,2,2,2)) if n>1 then for i = 1,nModules do b2:add(Residual(f,f)) end else for i = 1,nModules do b2:add(Residual(f,f)) end end if n>1 then b2:add(hourglass(n-1,f)) else for i = 1,nModules do b2:add(Residual(f,f)) end end if n>1 then for i = 1,nModules do b2:add(Residual(f,f)) end else for i=1,nModules do b2:add(Residual(f,f)) end end b2:add(upsample(2)) branch:add(b1):add(b2) model:add(branch) return model:add(nn.CAddTable()) end local function hourglass(n, f, inp) -- Upper branch local up1 = inp for i = 1,nModules do up1 = Residual(f,f)(up1) end -- Lower branch local low1 = cudnn.SpatialMaxPooling(2,2,2,2)(inp) for i = 1,nModules do low1 = Residual(f,f)(low1) end local low2 if n > 1 then low2 = hourglass(n-1,f,low1) else low2 = low1 for i = 1,nModules do low2 = Residual(f,f)(low2) end end local low3 = low2 for i = 1,nModules do low3 = Residual(f,f)(low3) end local up2 = nn.SpatialUpSamplingNearest(2)(low3) -- Bring two branches together return nn.CAddTable()({up1,up2}) end function createModel(opt) nModules = opt.nModules nFeats = opt.nFeats nStack = opt.nStacks local inp = nn.Identity()() -- Initial processing of the image local cnv1_ = conv(3,64,7,7,2,2,3,3)(inp) -- 128 local cnv1 = relu(true)(batchnorm(64)(cnv1_)) local r1 = Residual(64,128)(cnv1) local pool = nn.SpatialMaxPooling(2,2,2,2)(r1) -- 64 local r4 = Residual(128,128)(pool) local r5 = Residual(128,nFeats)(r4) local out = {} local inter = r5 for i = 1,nStack do local hg = hourglass(4,nFeats,inter) -- Residual layers at output resolution local ll = hg for j = 1,nModules do ll = Residual(nFeats,nFeats)(ll) end -- Linear layer to produce first set of predictions ll = lin(nFeats,nFeats,ll) -- Predicted heatmaps local tmpOut = conv(nFeats,68,1,1,1,1,0,0)(ll) table.insert(out,tmpOut) -- Add predictions back if i < nStack then local ll_ = conv(nFeats,nFeats,1,1,1,1,0,0)(ll) local tmpOut_ = conv(68,nFeats,1,1,1,1,0,0)(tmpOut) inter = nn.CAddTable()({inter, ll_, tmpOut_}) end end -- Final model local model = nn.gModule({inp}, out) return model end return createModel ================================================ FILE: models/init.lua ================================================ local M = {} function M.setup(opt, checkpoint) local model if checkpoint then local modelPath = paths.concat(opt.resume, checkpoint.modelFile) assert(paths.filep(modelPath), 'Saved model not found: ' .. modelPath) print('=> Resuming model from ' .. modelPath) model = torch.load(modelPath) elseif opt.retrain ~= 'none' then local modelPath = paths.concat(opt.resume, checkpoint.modelFile) assert(paths.filep(modelPath), 'Saved model not found: ' .. modelPath) print('=> Resuming model from ' .. modelPath) model = torch.load(modelPath) if preprocess == false then return model, nil end else print('=> Creating model from file: models/' .. opt.netType .. '.lua') model = require('models/' .. opt.netType)(opt) end if torch.type(model) == 'nn.DataParallelTable' then model = model:get(1) end -- Set the CUDNN flags if opt.cudnn == 'fastest' then cudnn.fastest = true cudnn.benchmark = true elseif opt.cudnn == 'deterministic' then -- Use a deterministic convolution implementation model:apply(function(m) if m.setMode then m:setMode(1, 1, 1) end end) end if opt.nGPU > 1 then local gpus = torch.range(1, opt.nGPU):totable() local fastest, benchmark = cudnn.fastest, cudnn.benchmark local dpt = nn.DataParallelTable(1, true, true) :add(model, gpus) :threads(function() local cudnn = require 'cudnn' require 'nngraph' require 'newLayers.BinActiveZ' cudnn.fastest, cudnn.benchmark = fastest, benchmark end) dpt.gradInput = nil model = dpt:cuda() end local criterion if opt.nStacks>1 then criterion = nn.ParallelCriterion() for i=1,opt.nStacks do criterion:add(nn.MSECriterion()) end else criterion = nn.MSECriterion() end return model:cuda(), criterion:cuda() end return M ================================================ FILE: opts.lua ================================================ local M = { } function M.parse(arg) local cmd = torch.CmdLine() cmd:text() cmd:text('2D-FAN and 3D-FAN Training script') cmd:text('Visit https://www.adrianbulat.com for more details') cmd:text() cmd:text('Options:') ------------ General options -------------------- cmd:option('-data', 'dataset/300W-LP/', 'Path to dataset') cmd:option('-manualSeed', 0, 'Manually set RNG seed') cmd:option('-nGPU', 1, 'Number of GPUs to use by default') cmd:option('-backend', 'cudnn', 'Options: cudnn | cunn') cmd:option('-cudnn', 'fastest', 'Options: fastest | default | deterministic') cmd:option('-gen', 'gen', 'Path to save generated files') cmd:option('-snapshot', 3, 'save a snapshot every n epochs') ------------- Data options ------------------------ cmd:option('-nThreads', 2, 'number of data loading threads') ------------- Training options -------------------- cmd:option('-nEpochs', 100, 'Number of total epochs to run') cmd:option('-epochNumber', 1, 'Manual epoch number (useful on restarts)') cmd:option('-batchSize', 10, 'mini-batch size (1 = pure stochastic)') ------------- Checkpointing options --------------- cmd:option('-save', 'checkpoints', 'Directory in which to save checkpoints') cmd:option('-resume', 'none', 'Resume from the latest checkpoint in this directory') ---------- Optimization options ---------------------- cmd:option('-LR', 0.00025, 'initial learning rate') cmd:option('-momentum', 0.0, 'momentum') cmd:option('-weightDecay', 0.0, 'weight decay') ---------- Model options ---------------------------------- cmd:option('-netType', 'fan', 'Options: fan') cmd:option('-nModules', 1, 'Number of modues per level') cmd:option('-nStacks', 4, 'Number of stacked networks') cmd:option('-nFeats', 256, 'BLock width (# channels)') cmd:option('-retrain', 'none', 'Path to model to retrain with') cmd:option('-optimState', 'none', 'Path to an optimState to reload from') ---------- Augumentation options ---------------------------------- cmd:option('-scaleFactor', 0.3, 'scaling factor') cmd:option('-rotFactor', 30, 'rotation factor (in degrees)') cmd:text() local opt = cmd:parse(arg or {}) if not paths.dirp(opt.save) and not paths.mkdir(opt.save) then cmd:error('error: unable to create checkpoint directory: ' .. opt.save .. '\n') end return opt end return M ================================================ FILE: train.lua ================================================ require 'cunn' local optim = require 'optim' local lr_policy = { {0,50,2.5e-4}, {50,70,1e-4}, {70,90,5e-5}, {90,100,1e-5}, {100,110,5e-6} } local M = {} local Trainer = torch.class('Trainer', M) function Trainer:__init(model,criterion,opt,optimState) self.model = model self.criterion = criterion self.optimState = optimState or { learningRate = opt.LR, learningRateDecay = 0.0, momentum = opt.momentum, epsilon = 1e-8, weightDecay = opt.weightDecay, } self.opt = opt self.params, self.gradParams = model:getParameters() end function Trainer:train(epoch, dataloader) local avgLoss, avgAcc = 0.0, 0.0 self.optimState.learningRate = self:learningRate(epoch) local timer = torch.Timer() local dataTimer = torch.Timer() local function feval() return self.criterion.output, self.gradParams end local trainSize = dataloader:size() local N = 0 print('=> Training epoch # '..epoch) self.model:training() for n, sample in dataloader:run() do local dataTime = dataTimer:time().real self:copyInputs(sample) self.model:zeroGradParameters() local output = self.model:forward(self.input) local loss = self.criterion:forward(output, self.label) self.criterion:backward(self.model.output, self.label) self.model:backward(self.input,self.criterion.gradInput) optim.rmsprop(feval, self.params, self.optimState) avgLoss = avgLoss + loss N = N + 1 print((' | Epoch: [%d][%d/%d] Time %.3f Data %.3f Err %1.4f'):format( epoch, n, trainSize, timer:time().real, dataTime, loss)) -- check that the storage didn't get changed do to an unfortunate getParameters call assert(self.params:storage() == self.model:parameters()[1]:storage()) collectgarbage() timer:reset() dataTimer:reset() end return avgLoss / N, avgAcc / N end function Trainer:learningRate(epoch) local decay = 0 for i=1, #lr_policy do if (epoch>lr_policy[i][1]) and (lr_policy[i][2]>=epoch) then print(string.format('Using lr_rate: %f',lr_policy[i][3])) return lr_policy[i][3] end end end function Trainer:copyInputs(sample) -- Copies the input to a CUDA tensor, if using 1 GPU, or to pinned memory, -- if using DataParallelTable. The target is always copied to a CUDA tensor self.input = self.input or (self.opt.nGPU == 1 and torch.CudaTensor() or cutorch.createCudaHostTensor()) label = label or torch.CudaTensor() self.input:resize(sample.input[{{},{},{},{}}]:size()):copy(sample.input[{{},{},{},{}}]) label:resize(sample.label:size()):copy(sample.label) -- Adjust the input accordingly to the network arhitecture if self.opt.nStacks>1 then local tempLabel = {} for i=1,self.opt.nStacks do table.insert(tempLabel, label) end self.label = tempLabel else self.label = label end end return M.Trainer ================================================ FILE: utils.lua ================================================ ------------------------------------------------------------------------------- -- Coordinate transformation ------------------------------------------------------------------------------- function getTransform(center, scale, rot, res) local h = 200 * scale local t = torch.eye(3) -- Scaling t[1][1] = res / h t[2][2] = res / h -- Translation t[1][3] = res * (-center[1] / h + .5) t[2][3] = res * (-center[2] / h + .5) -- Rotation if rot ~= 0 then rot = -rot local r = torch.eye(3) local ang = rot * math.pi / 180 local s = math.sin(ang) local c = math.cos(ang) r[1][1] = c r[1][2] = -s r[2][1] = s r[2][2] = c -- Need to make sure rotation is around center local t_ = torch.eye(3) t_[1][3] = -res/2 t_[2][3] = -res/2 local t_inv = torch.eye(3) t_inv[1][3] = res/2 t_inv[2][3] = res/2 t = t_inv * r * t_ * t end return t end function transform(pt, center, scale, rot, res, invert) local pt_ = torch.ones(3) pt_[1],pt_[2] = pt[1]-1,pt[2]-1 local t = getTransform(center, scale, rot, res) if invert then t = torch.inverse(t) end local new_point = (t*pt_):sub(1,2):add(1e-4) return new_point:int():add(1) end function crop(img, center, scale, rot, res) local ul = transform({1,1}, center, scale, 0, res, true) local br = transform({res+1,res+1}, center, scale, 0, res, true) local pad = math.floor(torch.norm((ul - br):float())/2 - (br[1]-ul[1])/2) if rot ~= 0 then ul = ul - pad br = br + pad end local newDim,newImg,ht,wd if img:size():size() > 2 then newDim = torch.IntTensor({img:size(1), br[2] - ul[2], br[1] - ul[1]}) newImg = torch.zeros(newDim[1],newDim[2],newDim[3]) ht = img:size(2) wd = img:size(3) else newDim = torch.IntTensor({br[2] - ul[2], br[1] - ul[1]}) newImg = torch.zeros(newDim[1],newDim[2]) ht = img:size(1) wd = img:size(2) end local newX = torch.Tensor({math.max(1, -ul[1] + 2), math.min(br[1], wd+1) - ul[1]}) local newY = torch.Tensor({math.max(1, -ul[2] + 2), math.min(br[2], ht+1) - ul[2]}) local oldX = torch.Tensor({math.max(1, ul[1]), math.min(br[1], wd+1) - 1}) local oldY = torch.Tensor({math.max(1, ul[2]), math.min(br[2], ht+1) - 1}) if newDim:size(1) > 2 then newImg:sub(1,newDim[1],newY[1],newY[2],newX[1],newX[2]):copy(img:sub(1,newDim[1],oldY[1],oldY[2],oldX[1],oldX[2])) else newImg:sub(newY[1],newY[2],newX[1],newX[2]):copy(img:sub(oldY[1],oldY[2],oldX[1],oldX[2])) end if rot ~= 0 then newImg = image.rotate(newImg, rot * math.pi / 180, 'bilinear') if newDim:size(1) > 2 then newImg = newImg:sub(1,newDim[1],pad,newDim[2]-pad,pad,newDim[3]-pad) else newImg = newImg:sub(pad,newDim[1]-pad,pad,newDim[2]-pad) end end newImg = image.scale(newImg,res,res) return newImg end local magic_gaussian = image.gaussian(7) function drawGaussian(img, pt, sigma) -- Check if the gaussian is in-bounds local ul = {math.floor(pt[1] - 3 * sigma), math.floor(pt[2] - 3 * sigma)} local br = {math.floor(pt[1] + 3 * sigma), math.floor(pt[2] + 3 * sigma)} -- return the image otherwise if (ul[1] > img:size(2) or ul[2] > img:size(1) or br[1] < 1 or br[2] < 1) then return img end -- Generate gaussian local size = 6 * sigma + 1 -- Avoid the need of generating the gaussian for each sample local g = magic_gaussian:clone()--image.gaussian(size) -- , 1 / size, 1) -- Usable gaussian range local g_x = {math.max(1, -ul[1]), math.min(br[1], img:size(2)) - math.max(1, ul[1]) + math.max(1, -ul[1])} local g_y = {math.max(1, -ul[2]), math.min(br[2], img:size(1)) - math.max(1, ul[2]) + math.max(1, -ul[2])} -- Image range local img_x = {math.max(1, ul[1]), math.min(br[1], img:size(2))} local img_y = {math.max(1, ul[2]), math.min(br[2], img:size(1))} assert(g_x[1] > 0 and g_y[1] > 0) img:sub(img_y[1], img_y[2], img_x[1], img_x[2]):add(g:sub(g_y[1], g_y[2], g_x[1], g_x[2])) img[img:gt(1)] = 1 return img end function shuffleLR(x) local dim if x:nDimension() == 4 then dim = 2 else assert(x:nDimension() == 3) dim = 1 end -- Keypoints pairs for 300W_LP, 300VW, 300W and LS3D-W datasets local matchedParts = { {1,17}, {2,16}, {3,15}, {4,14}, {5,13}, {6,12}, {7,11}, {8,10}, {18,27},{19,26},{20,25},{21,24},{22,23}, {37,46},{38,45},{39,44},{40,43}, {42,47},{41,48}, {32,36},{33,35}, {51,53},{50,54},{49,55},{62,64},{61,65},{68,66},{60,56}, {59,57} } for i = 1,#matchedParts do local idx1, idx2 = unpack(matchedParts[i]) local tmp = x:narrow(dim, idx1, 1):clone() x:narrow(dim, idx1, 1):copy(x:narrow(dim, idx2, 1)) x:narrow(dim, idx2, 1):copy(tmp) end return x end function flip(x) require 'image' local y = torch.FloatTensor(x:size()) for i = 1, x:size(1) do image.hflip(y[i], x[i]:float()) end return y:typeAs(x) end