229 lines
7.5 KiB
Lua
229 lines
7.5 KiB
Lua
local class = require 'class'
|
|
require 'models.base_model'
|
|
require 'models.architectures'
|
|
require 'util.image_pool'
|
|
util = paths.dofile('../util/util.lua')
|
|
Pix2PixModel = class('Pix2PixModel', 'BaseModel')
|
|
|
|
function Pix2PixModel:__init(conf)
|
|
conf = conf or {}
|
|
end
|
|
|
|
-- Returns the name of the model
|
|
function Pix2PixModel:model_name()
|
|
return 'Pix2PixModel'
|
|
end
|
|
|
|
function Pix2PixModel:InitializeStates()
|
|
return {learningRate=opt.lr, beta1=opt.beta1,}
|
|
end
|
|
|
|
-- Defines models and networks
|
|
function Pix2PixModel:Initialize(opt) -- use lsgan
|
|
-- define tensors
|
|
local d_input_nc = opt.input_nc + opt.output_nc
|
|
self.real_AB = torch.Tensor(opt.batchSize, d_input_nc, opt.fineSize, opt.fineSize)
|
|
self.fake_AB = torch.Tensor(opt.batchSize, d_input_nc, opt.fineSize, opt.fineSize)
|
|
if opt.test == 0 then
|
|
self.fakeABPool = ImagePool(opt.pool_size)
|
|
end
|
|
-- load/define models
|
|
self.criterionGAN = nn.MSECriterion()
|
|
self.criterionL1 = nn.AbsCriterion()
|
|
|
|
local netG, netD = nil, nil
|
|
if opt.continue_train == 1 then
|
|
if opt.test == 1 then -- only load model G for test
|
|
netG = util.load_test_model('G', opt)
|
|
else
|
|
netG = util.load_model('G', opt)
|
|
netD = util.load_model('D', opt)
|
|
end
|
|
else
|
|
netG = defineG(opt.input_nc, opt.output_nc, opt.ngf, opt.which_model_netG)
|
|
netD = defineD(d_input_nc, opt.ndf, opt.which_model_netD, opt.n_layers_D, false) -- with sigmoid
|
|
end
|
|
|
|
self.netD = netD
|
|
self.netG = netG
|
|
|
|
-- define real/fake labels
|
|
if opt.test == 0 then
|
|
netD_output_size = self.netD:forward(self.real_AB):size()
|
|
self.fake_label = torch.Tensor(netD_output_size):fill(0.0)
|
|
self.real_label = torch.Tensor(netD_output_size):fill(1.0) -- no soft smoothing
|
|
|
|
self.optimStateD = self:InitializeStates()
|
|
self.optimStateG = self:InitializeStates()
|
|
|
|
self:RefreshParameters()
|
|
|
|
print('---------- # Learnable Parameters --------------')
|
|
print(('G = %d'):format(self.parametersG:size(1)))
|
|
print(('D = %d'):format(self.parametersD:size(1)))
|
|
print('------------------------------------------------')
|
|
end
|
|
|
|
self.A_idx = {{}, {1, opt.input_nc}, {}, {}}
|
|
self.B_idx = {{}, {opt.input_nc+1, opt.input_nc+opt.output_nc}, {}, {}}
|
|
end
|
|
|
|
-- Runs the forward pass of the network
|
|
function Pix2PixModel:Forward(input, opt)
|
|
if opt.which_direction == 'BtoA' then
|
|
local temp = input.real_A
|
|
input.real_A = input.real_B
|
|
input.real_B = temp
|
|
end
|
|
|
|
if opt.test == 0 then
|
|
self.real_AB[self.A_idx]:copy(input.real_A)
|
|
self.real_AB[self.B_idx]:copy(input.real_B)
|
|
self.real_A = self.real_AB[self.A_idx]
|
|
self.real_B = self.real_AB[self.B_idx]
|
|
|
|
self.fake_AB[self.A_idx]:copy(self.real_A)
|
|
self.fake_B = self.netG:forward(self.real_A):clone()
|
|
self.fake_AB[self.B_idx]:copy(self.fake_B)
|
|
else
|
|
if opt.gpu > 0 then
|
|
self.real_A = input.real_A:cuda()
|
|
self.real_B = input.real_B:cuda()
|
|
else
|
|
self.real_A = input.real_A:clone()
|
|
self.real_B = input.real_B:clone()
|
|
end
|
|
self.fake_B = self.netG:forward(self.real_A):clone()
|
|
end
|
|
end
|
|
|
|
-- create closure to evaluate f(X) and df/dX of discriminator
|
|
function Pix2PixModel:fDx_basic(x, gradParams, netD, netG, real, fake, opt)
|
|
util.BiasZero(netD)
|
|
util.BiasZero(netG)
|
|
gradParams:zero()
|
|
|
|
-- Real log(D(B))
|
|
local output = netD:forward(real)
|
|
local errD_real = self.criterionGAN:forward(output, self.real_label)
|
|
local df_do = self.criterionGAN:backward(output, self.real_label)
|
|
netD:backward(real, df_do)
|
|
-- Fake + log(1 - D(G(A)))
|
|
output = netD:forward(fake)
|
|
local errD_fake = self.criterionGAN:forward(output, self.fake_label)
|
|
local df_do2 = self.criterionGAN:backward(output, self.fake_label)
|
|
netD:backward(fake, df_do2)
|
|
-- calculate loss
|
|
local errD = (errD_real + errD_fake) / 2.0
|
|
return errD, gradParams
|
|
end
|
|
|
|
|
|
function Pix2PixModel:fDx(x, opt)
|
|
fake_AB = self.fakeABPool:Query(self.fake_AB)
|
|
self.errD, gradParams = self:fDx_basic(x, self.gradParametersD, self.netD, self.netG,
|
|
self.real_AB, fake_AB, opt)
|
|
return self.errD, gradParams
|
|
end
|
|
|
|
function Pix2PixModel:fGx_basic(x, netG, netD, real, fake, gradParametersG, opt)
|
|
util.BiasZero(netG)
|
|
util.BiasZero(netD)
|
|
gradParametersG:zero()
|
|
|
|
-- First. G(A) should fake the discriminator
|
|
local output = netD:forward(fake)
|
|
local errG = self.criterionGAN:forward(output, self.real_label)
|
|
local dgan_loss_dd = self.criterionGAN:backward(output, self.real_label)
|
|
local dgan_loss_do = netD:updateGradInput(fake, dgan_loss_dd)
|
|
|
|
-- Second. G(A) should be close to the real
|
|
real_B = real[self.B_idx]
|
|
real_A = real[self.A_idx]
|
|
fake_B = fake[self.B_idx]
|
|
local errL1 = self.criterionL1:forward(fake_B, real_B) * opt.lambda_A
|
|
local dl1_loss_do = self.criterionL1:backward(fake_B, real_B) * opt.lambda_A
|
|
netG:backward(real_A, dgan_loss_do[self.B_idx] + dl1_loss_do)
|
|
|
|
return gradParametersG, errG, errL1
|
|
end
|
|
|
|
function Pix2PixModel:fGx(x, opt)
|
|
self.gradParametersG, self.errG, self.errL1 = self:fGx_basic(x, self.netG, self.netD,
|
|
self.real_AB, self.fake_AB, self.gradParametersG, opt)
|
|
return self.errG, self.gradParametersG
|
|
end
|
|
|
|
-- Runs the backprop gradient descent
|
|
-- Corresponds to a single batch of data
|
|
function Pix2PixModel:OptimizeParameters(opt)
|
|
local fD = function(x) return self:fDx(x, opt) end
|
|
local fG = function(x) return self:fGx(x, opt) end
|
|
optim.adam(fD, self.parametersD, self.optimStateD)
|
|
optim.adam(fG, self.parametersG, self.optimStateG)
|
|
end
|
|
|
|
-- This function can be used to reset momentum after each epoch
|
|
function Pix2PixModel:RefreshParameters()
|
|
self.parametersD, self.gradParametersD = nil, nil -- nil them to avoid spiking memory
|
|
self.parametersG, self.gradParametersG = nil, nil
|
|
|
|
-- define parameters of optimization
|
|
self.parametersG, self.gradParametersG = self.netG:getParameters()
|
|
self.parametersD, self.gradParametersD = self.netD:getParameters()
|
|
end
|
|
|
|
-- This function updates the learning rate; lr for the first opt.niter iterations; graduatlly decreases the lr to 0 for the next opt.niter_decay iterations
|
|
function Pix2PixModel:UpdateLearningRate(opt)
|
|
local lrd = opt.lr / opt.niter_decay
|
|
local old_lr = self.optimStateD['learningRate']
|
|
local lr = old_lr - lrd
|
|
self.optimStateD['learningRate'] = lr
|
|
self.optimStateG['learningRate'] = lr
|
|
print(('update learning rate: %f -> %f'):format(old_lr, lr))
|
|
end
|
|
|
|
|
|
-- Save the current model to the file system
|
|
function Pix2PixModel:Save(prefix, opt)
|
|
util.save_model(self.netG, prefix .. '_net_G.t7', 1.0)
|
|
util.save_model(self.netD, prefix .. '_net_D.t7', 1.0)
|
|
end
|
|
|
|
-- returns a string that describes the current errors
|
|
function Pix2PixModel:GetCurrentErrorDescription()
|
|
description = ('G: %.4f D: %.4f L1: %.4f'):format(
|
|
self.errG and self.errG or -1, self.errD and self.errD or -1, self.errL1 and self.errL1 or -1)
|
|
return description
|
|
|
|
end
|
|
|
|
|
|
-- returns a string that describes the display plot configuration
|
|
function Pix2PixModel:DisplayPlot(opt)
|
|
return 'errG,errD,errL1'
|
|
end
|
|
|
|
|
|
-- returns current errors
|
|
function Pix2PixModel:GetCurrentErrors()
|
|
local errors = {errG=self.errG, errD=self.errD, errL1=self.errL1}
|
|
return errors
|
|
end
|
|
|
|
-- returns a table of image/label pairs that describe
|
|
-- the current results.
|
|
-- |return|: a table of table. List of image/label pairs
|
|
function Pix2PixModel:GetCurrentVisuals(opt, size)
|
|
if not size then
|
|
size = opt.display_winsize
|
|
end
|
|
|
|
local visuals = {}
|
|
table.insert(visuals, {img=self.real_A, label='real_A'})
|
|
table.insert(visuals, {img=self.fake_B, label='fake_B'})
|
|
table.insert(visuals, {img=self.real_B, label='real_B'})
|
|
|
|
return visuals
|
|
end
|