+
+-- Francois Fleuret's Torch Toolbox
+
+require 'torch'
+require 'nn'
+
+----------------------------------------------------------------------
+
+colors = sys.COLORS
+
+function printf(f, ...)
+ print(string.format(f, unpack({...})))
+end
+
+function printfc(c, f, ...)
+ printf(c .. string.format(f, unpack({...})) .. colors.black)
+end
+
+function logCommand(c)
+ print(colors.blue .. '[' .. c .. '] -> [' .. sys.execute(c) .. ']' .. colors.black)
+end
+
+----------------------------------------------------------------------
+-- Environment variables
+
+defaultNbThreads = 1
+defaultUseGPU = false
+
+if os.getenv('TORCH_NB_THREADS') then
+ defaultNbThreads = os.getenv('TORCH_NB_THREADS')
+ print('Environment variable TORCH_NB_THREADS is set and equal to ' .. defaultNbThreads)
+else
+ print('Environment variable TORCH_NB_THREADS is not set, default is ' .. defaultNbThreads)
+end
+
+if os.getenv('TORCH_USE_GPU') then
+ defaultUseGPU = os.getenv('TORCH_USE_GPU') == 'yes'
+ print('Environment variable TORCH_USE_GPU is set and evaluated as ' .. tostring(defaultUseGPU))
+else
+ print('Environment variable TORCH_USE_GPU is not set, default is ' .. tostring(defaultUseGPU))
+end
+
+----------------------------------------------------------------------
+
+function fftbInit(cmd, params)
+
+ torch.setnumthreads(params.nbThreads)
+ torch.setdefaulttensortype('torch.FloatTensor')
+ torch.manualSeed(params.seed)
+
+ -- Logging
+
+ if params.rundir == '' then
+ params.rundir = cmd:string('experiment', params, { })
+ end
+
+ paths.mkdir(params.rundir)
+
+ if not params.noLog then
+ -- Append to the log if there is one
+ cmd:log(io.open(params.rundir .. '/log', 'a'), params)
+ end
+
+ -- Dealing with the CPU/GPU
+
+ ffnn = {}
+
+ -- By default, ffnn returns the entries from nn
+ local mt = {}
+ function mt.__index(table, key)
+ return (cudnn and cudnn[key]) or (cunn and cunn[key]) or nn[key]
+ end
+ setmetatable(ffnn, mt)
+
+ -- These are the tensors that can be kept on the CPU
+ ffnn.SlowTensor = torch.Tensor
+ ffnn.SlowStorage = torch.Storage
+ -- These are the tensors that should be moved to the GPU
+ ffnn.FastTensor = torch.Tensor
+ ffnn.FastStorage = torch.Storage
+
+ if params.useGPU then
+ require 'cutorch'
+ require 'cunn'
+ require 'cudnn'
+
+ if params.fastGPU then
+ cudnn.benchmark = true
+ cudnn.fastest = true
+ end
+
+ ffnn.FastTensor = torch.CudaTensor
+ ffnn.FastStorage = torch.CudaStorage
+ end
+end
+
+----------------------------------------------------------------------
+
+function dimAtThatPoint(model, input)
+ if params.useGPU then
+ model:cuda()
+ end
+ local i = ffnn.FastTensor(input:narrow(1, 1, 1):size()):copy(input:narrow(1, 1, 1))
+ return model:forward(i):nElement()
+end
+
+----------------------------------------------------------------------
+
+function sizeForBatch(n, x)
+ local size = x:size()
+ size[1] = n
+ return size
+end
+
+function fillBatch(data, first, batch, permutation)
+ local actualBatchSize = math.min(params.batchSize, data.input:size(1) - first + 1)
+
+ if batch.input then
+ if actualBatchSize ~= batch.input:size(1) then
+ batch.input:resize(sizeForBatch(actualBatchSize, batch.input))
+ end
+ else
+ if torch.isTypeOf(data.input, ffnn.SlowTensor) then
+ batch.input = ffnn.FastTensor(sizeForBatch(actualBatchSize, data.input));
+ else
+ batch.input = data.input.new():resize(sizeForBatch(actualBatchSize, data.input));
+ end
+ end
+
+ if batch.target then
+ if actualBatchSize ~= batch.target:size(1) then
+ batch.target:resize(sizeForBatch(actualBatchSize, batch.target))
+ end
+ else
+ if torch.isTypeOf(data.target, ffnn.SlowTensor) then
+ batch.target = ffnn.FastTensor(sizeForBatch(actualBatchSize, data.target));
+ else
+ batch.target = data.target.new():resize(sizeForBatch(actualBatchSize, data.target));
+ end
+ end
+
+ for k = 1, actualBatchSize do
+ local i
+ if permutation then
+ i = permutation[first + k - 1]
+ else
+ i = first + k - 1
+ end
+ batch.input[k] = data.input[i]
+ batch.target[k] = data.target[i]
+ end
+end
+
+----------------------------------------------------------------------
+
+--[[
+
+The combineImage function takes as input a parameter c which is the
+value to use for the background of the resulting image (padding and
+such), and t which is either a 2d tensor, a 3d tensor, or a table.
+
+ * If t is a 3d tensor, it is returned unchanged.
+
+ * If t is a 2d tensor [r x c], it is reshaped to [1 x r x c] and
+ returned.
+
+ * If t is a table, combineImage first calls itself recursively on
+ t[1], t[2], etc.
+
+ It then creates a new tensor by concatenating the results
+ horizontally if t.vertical is nil, vertically otherwise.
+
+ It adds a padding of t.pad pixels if this field is set.
+
+ * Example
+
+ x = torch.Tensor(64, 64):fill(0.5)
+ y = torch.Tensor(100, 30):fill(0.85)
+
+ i = combineImages(1.0,
+ {
+ pad = 1,
+ vertical = true,
+ { pad = 1, x },
+ {
+ y,
+ { pad = 4, torch.Tensor(32, 16):fill(0.25) },
+ { pad = 1, torch.Tensor(45, 54):uniform(0.25, 0.9) },
+ }
+ }
+ )
+
+ image.save('example.png', i)
+
+]]--
+
+function combineImages(c, t)
+
+ if torch.isTensor(t) then
+
+ if t:dim() == 3 then
+ return t
+ elseif t:dim() == 2 then
+ return torch.Tensor(1, t:size(1), t:size(2)):copy(t)
+ else
+ error('can only deal with [height x width] or [channel x height x width] tensors.')
+ end
+
+ else
+
+ local subImages = {} -- The subimages
+ local nc = 0 -- Nb of columns
+ local nr = 0 -- Nb of rows
+
+ for i, x in ipairs(t) do
+ subImages[i] = combineImages(c, x)
+ if t.vertical then
+ nr = nr + subImages[i]:size(2)
+ nc = math.max(nc, subImages[i]:size(3))
+ else
+ nr = math.max(nr, subImages[i]:size(2))
+ nc = nc + subImages[i]:size(3)
+ end
+ end
+
+ local pad = t.pad or 0
+ local result = torch.Tensor(subImages[1]:size(1), nr + 2 * pad, nc + 2 * pad):fill(c)
+ local co = 1 + pad -- Origin column
+ local ro = 1 + pad -- Origin row
+
+ for i in ipairs(t) do
+
+ result
+ :sub(1, subImages[1]:size(1),
+ ro, ro + subImages[i]:size(2) - 1,
+ co, co + subImages[i]:size(3) - 1)
+ :copy(subImages[i])
+
+ if t.vertical then
+ ro = ro + subImages[i]:size(2)
+ else
+ co = co + subImages[i]:size(3)
+ end
+
+ end
+
+ return result
+
+ end
+
+end
+
+--[[
+
+The imageFromTensors function gets as input a list of tensors of
+arbitrary dimensions each, but whose two last dimensions stand for
+height x width. It creates an image tensor (2d, one channel) with each
+argument tensor unfolded per row.
+
+]]--
+
+function imageFromTensors(bt, signed)
+ local gap = 1
+ local tgap = -1
+ local width = 0
+ local height = gap
+
+ for _, t in pairs(bt) do
+ local d = t:dim()
+ local h, w = t:size(d - 1), t:size(d)
+ local n = t:nElement() / (w * h)
+ width = math.max(width, gap + n * (gap + w))
+ height = height + gap + tgap + gap + h
+ end
+
+ local e = torch.Tensor(3, height, width):fill(1.0)
+ local y0 = 1 + gap
+
+ for _, t in pairs(bt) do
+ local d = t:dim()
+ local h, w = t:size(d - 1), t:size(d)
+ local n = t:nElement() / (w * h)
+ local z = t:norm() / math.sqrt(t:nElement())
+
+ local x0 = 1 + gap + math.floor( (width - n * (w + gap)) /2 )
+ local u = torch.Tensor(t:size()):copy(t):resize(n, h, w)
+ for m = 1, n do
+
+ for c = 1, 3 do
+ for y = 0, h+1 do
+ e[c][y0 + y - 1][x0 - 1] = 0.0
+ e[c][y0 + y - 1][x0 + w ] = 0.0
+ end
+ for x = 0, w+1 do
+ e[c][y0 - 1][x0 + x - 1] = 0.0
+ e[c][y0 + h ][x0 + x - 1] = 0.0
+ end
+ end
+
+ for y = 1, h do
+ for x = 1, w do
+ local v = u[m][y][x] / z
+ local r, g, b
+ if signed then
+ if v < -1 then
+ r, g, b = 0.0, 0.0, 1.0
+ elseif v > 1 then
+ r, g, b = 1.0, 0.0, 0.0
+ elseif v >= 0 then
+ r, g, b = 1.0, 1.0 - v, 1.0 - v
+ else
+ r, g, b = 1.0 + v, 1.0 + v, 1.0
+ end
+ else
+ if v <= 0 then
+ r, g, b = 1.0, 1.0, 1.0
+ elseif v > 1 then
+ r, g, b = 0.0, 0.0, 0.0
+ else
+ r, g, b = 1.0 - v, 1.0 - v, 1.0 - v
+ end
+ end
+ e[1][y0 + y - 1][x0 + x - 1] = r
+ e[2][y0 + y - 1][x0 + x - 1] = g
+ e[3][y0 + y - 1][x0 + x - 1] = b
+ end
+ end
+ x0 = x0 + w + gap
+ end
+ y0 = y0 + h + gap + tgap + gap
+ end
+
+ return e
+end