require './lib/portable' require './lib/mynn' require 'optim' require 'xlua' require 'pl' require 'snappy' local settings = require './lib/settings' local srcnn = require './lib/srcnn' local minibatch_adam = require './lib/minibatch_adam' local iproc = require './lib/iproc' local reconstruct = require './lib/reconstruct' local pairwise_transform = require './lib/pairwise_transform' local image_loader = require './lib/image_loader' local function save_test_scale(model, rgb, file) local up = reconstruct.scale(model, settings.scale, rgb) image.save(file, up) end local function save_test_jpeg(model, rgb, file) local im, count = reconstruct.image(model, rgb) image.save(file, im) end local function split_data(x, test_size) local index = torch.randperm(#x) local train_size = #x - test_size local train_x = {} local valid_x = {} for i = 1, train_size do train_x[i] = x[index[i]] end for i = 1, test_size do valid_x[i] = x[index[train_size + i]] end return train_x, valid_x end local function make_validation_set(x, transformer, n) n = n or 4 local data = {} for i = 1, #x do for k = 1, math.max(n / 8, 1) do local xy = transformer(x[i], true, 8) for j = 1, #xy do local x = xy[j][1] local y = xy[j][2] table.insert(data, {x = x:reshape(1, x:size(1), x:size(2), x:size(3)), y = y:reshape(1, y:size(1), y:size(2), y:size(3))}) end end xlua.progress(i, #x) collectgarbage() end return data end local function validate(model, criterion, data) local loss = 0 for i = 1, #data do local z = model:forward(data[i].x:cuda()) loss = loss + criterion:forward(z, data[i].y:cuda()) xlua.progress(i, #data) if i % 10 == 0 then collectgarbage() end end return loss / #data end local function create_criterion(model) if reconstruct.is_rgb(model) then local offset = reconstruct.offset_size(model) local output_w = settings.crop_size - offset * 2 local weight = torch.Tensor(3, output_w * output_w) weight[1]:fill(0.299 * 3) -- R weight[2]:fill(0.587 * 3) -- G weight[3]:fill(0.114 * 3) -- B return mynn.RGBWeightedMSECriterion(weight):cuda() else return nn.MSECriterion():cuda() end end local function transformer(x, is_validation, n, offset) local size = x[1] local dec = snappy.decompress(x[2]:string()) x = torch.ByteTensor(size[1], size[2], size[3]) x:storage():string(dec) n = n or settings.batch_size; if is_validation == nil then is_validation = false end local color_noise = nil local overlay = nil local active_cropping_ratio = nil local active_cropping_tries = nil if is_validation then active_cropping_rate = 0.0 active_cropping_tries = 0 color_noise = false overlay = false else active_cropping_rate = settings.active_cropping_rate active_cropping_tries = settings.active_cropping_tries color_noise = settings.color_noise overlay = settings.overlay end if settings.method == "scale" then return pairwise_transform.scale(x, settings.scale, settings.crop_size, offset, n, { color_noise = color_noise, overlay = overlay, random_half = settings.random_half, active_cropping_rate = active_cropping_rate, active_cropping_tries = active_cropping_tries, rgb = (settings.color == "rgb") }) elseif settings.method == "noise" then return pairwise_transform.jpeg(x, settings.category, settings.noise_level, settings.crop_size, offset, n, { color_noise = color_noise, overlay = overlay, active_cropping_rate = active_cropping_rate, active_cropping_tries = active_cropping_tries, random_half = settings.random_half, jpeg_sampling_factors = settings.jpeg_sampling_factors, rgb = (settings.color == "rgb") }) elseif settings.method == "noise_scale" then return pairwise_transform.jpeg_scale(x, settings.scale, settings.category, settings.noise_level, settings.crop_size, offset, n, { color_noise = color_noise, overlay = overlay, jpeg_sampling_factors = settings.jpeg_sampling_factors, random_half = settings.random_half, rgb = (settings.color == "rgb") }) end end local function train() local model = srcnn.create(settings.method, settings.backend, settings.color) local offset = reconstruct.offset_size(model) local pairwise_func = function(x, is_validation, n) return transformer(x, is_validation, n, offset) end local criterion = create_criterion(model) local x = torch.load(settings.images) local lrd_count = 0 local train_x, valid_x = split_data(x, math.floor(settings.validation_ratio * #x)) local adam_config = { learningRate = settings.learning_rate, xBatchSize = settings.batch_size, } local ch = nil if settings.color == "y" then ch = 1 elseif settings.color == "rgb" then ch = 3 end local best_score = 100000.0 print("# make validation-set") local valid_xy = make_validation_set(valid_x, pairwise_func, settings.validation_crops) valid_x = nil collectgarbage() model:cuda() print("load .. " .. #train_x) for epoch = 1, settings.epoch do model:training() print("# " .. epoch) print(minibatch_adam(model, criterion, train_x, adam_config, pairwise_func, {ch, settings.crop_size, settings.crop_size}, {ch, settings.crop_size - offset * 2, settings.crop_size - offset * 2} )) model:evaluate() print("# validation") local score = validate(model, criterion, valid_xy) if score < best_score then local test_image = image_loader.load_float(settings.test) -- reload lrd_count = 0 best_score = score print("* update best model") torch.save(settings.model_file, model) if settings.method == "noise" then local log = path.join(settings.model_dir, ("noise%d_best.png"):format(settings.noise_level)) save_test_jpeg(model, test_image, log) elseif settings.method == "scale" then local log = path.join(settings.model_dir, ("scale%.1f_best.png"):format(settings.scale)) save_test_scale(model, test_image, log) elseif settings.method == "noise_scale" then local log = path.join(settings.model_dir, ("noise%d_scale%.1f_best.png"):format(settings.noise_level, settings.scale)) save_test_scale(model, test_image, log) end else lrd_count = lrd_count + 1 if lrd_count > 5 then lrd_count = 0 adam_config.learningRate = adam_config.learningRate * 0.9 print("* learning rate decay: " .. adam_config.learningRate) end end print("current: " .. score .. ", best: " .. best_score) collectgarbage() end end torch.manualSeed(settings.seed) cutorch.manualSeed(settings.seed) print(settings) train()