def addentry(self, step, score, train_losses):
if step % 100 == 0:
test_score = self.ema_helper.ema_copy(score, self.args.device)
test_score.eval()
try:
test_X, _ = next(self.test_iter)
except StopIteration:
self.test_iter = iter(self.testloader)
test_X, _ = next(self.test_iter)
test_X = test_X.to(self.args.device)
test_X = data_transform(self.config.data, test_X)
if self.hook is not None:
self.hook.update_step(step)
test_dsm_loss, _, _ = anneal_dsm_score_estimation(self.args, test_score, test_X, self.sigmas, hook=self.hook)
if self.args.adversarial:
self.tb_logger.add_scalar('loss_D', train_losses[0], global_step=step)
self.tb_logger.add_scalar('loss_G_DAE', train_losses[1], global_step=step)
self.tb_logger.add_scalar('loss_G_GAN', train_losses[2], global_step=step)
else:
self.tb_logger.add_scalar('train_loss', train_losses[0], global_step=step)
self.tb_logger.add_scalar('test_loss', test_dsm_loss, global_step=step)
if self.args.adversarial:
_train_loss_string = "tloss_D: {:10.4f} | tloss_dae: {:10.4f} | tloss_adv: {:10.4f}"
else:
_train_loss_string = "train_loss: {:10.4f}"
string = "step: {:8d} | " + _train_loss_string + " || test_loss: {:10.4f}"
logging.info(string.format(step, *train_losses, test_dsm_loss.item()))
print(string.format(step, *train_losses, test_dsm_loss.item()))
def test(self):
score = get_model(self.config)
score = torch.nn.DataParallel(score)
sigmas = get_sigmas(self.config)
dataset, test_dataset = get_dataset(self.args, self.config)
test_dataloader = DataLoader(test_dataset,
batch_size=self.config.test.batch_size,
shuffle=True,
num_workers=self.config.data.num_workers,
drop_last=True)
verbose = False
for ckpt in tqdm.tqdm(range(self.config.test.begin_ckpt,
self.config.test.end_ckpt + 1, 5000),
desc="processing ckpt:"):
states = torch.load(os.path.join(self.args.log_path,
f'checkpoint_{ckpt}.pth'),
map_location=self.config.device)
if self.config.model.ema:
ema_helper = EMAHelper(mu=self.config.model.ema_rate)
ema_helper.register(score)
ema_helper.load_state_dict(states[-1])
ema_helper.ema(score)
else:
score.load_state_dict(states[0])
score.eval()
step = 0
mean_loss = 0.
mean_grad_norm = 0.
average_grad_scale = 0.
for x, y in test_dataloader:
step += 1
x = x.to(self.config.device)
x = data_transform(self.config, x)
with torch.no_grad():
test_loss = anneal_dsm_score_estimation(
score, x, sigmas, None,
self.config.training.anneal_power)
if verbose:
logging.info("step: {}, test_loss: {}".format(
step, test_loss.item()))
mean_loss += test_loss.item()
mean_loss /= step
mean_grad_norm /= step
average_grad_scale /= step
logging.info("ckpt: {}, average test loss: {}".format(
ckpt, mean_loss))
def train(self):
if self.config.data.random_flip is False:
tran_transform = test_transform = transforms.Compose([
transforms.Resize(self.config.data.image_size),
transforms.ToTensor()
])
else:
tran_transform = transforms.Compose([
transforms.Resize(self.config.data.image_size),
transforms.RandomHorizontalFlip(p=0.5),
transforms.ToTensor()
])
test_transform = transforms.Compose([
transforms.Resize(self.config.data.image_size),
transforms.ToTensor()
])
dataset = LoadDataset('./ground turth',tran_transform)
dataloader = DataLoader(dataset, batch_size=self.config.training.batch_size, shuffle=True, num_workers=4)
self.config.input_dim = self.config.data.image_size ** 2 * self.config.data.channels
tb_path = os.path.join(self.args.run, 'tensorboard', self.args.doc)
if os.path.exists(tb_path):
shutil.rmtree(tb_path)
tb_logger = tensorboardX.SummaryWriter(log_dir=tb_path)
score = CondRefineNetDilated(self.config).to(self.config.device)
score = torch.nn.DataParallel(score)
optimizer = self.get_optimizer(score.parameters())
if self.args.resume_training:
states = torch.load(os.path.join(self.args.log, 'checkpoint.pth'))
score.load_state_dict(states[0])
optimizer.load_state_dict(states[1])
step = 0
sigmas = torch.tensor(
np.exp(np.linspace(np.log(self.config.model.sigma_begin), np.log(self.config.model.sigma_end),
self.config.model.num_classes))).float().to(self.config.device)
for epoch in range(self.config.training.n_epochs):
for i, X in enumerate(dataloader):
X = torch.tensor(X,dtype=torch.float32)
step += 1
score.train()
X = X.to(self.config.device)
X = X / 256. * 255. + torch.rand_like(X) / 256.
if self.config.data.logit_transform:
X = self.logit_transform(X)
labels = torch.randint(0, len(sigmas), (X.shape[0],), device=X.device)
if self.config.training.algo == 'dsm':
loss = anneal_dsm_score_estimation(score, X, labels, sigmas, self.config.training.anneal_power)
elif self.config.training.algo == 'ssm':
loss = anneal_sliced_score_estimation_vr(score, X, labels, sigmas,
n_particles=self.config.training.n_particles)
optimizer.zero_grad()
loss.backward()
optimizer.step()
tb_logger.add_scalar('loss', loss, global_step=step)
logging.info("step: {}, loss: {}".format(step, loss.item()))
if step >= self.config.training.n_iters:
return 0
if step % self.config.training.snapshot_freq == 0:
states = [
score.state_dict(),
optimizer.state_dict(),
]
torch.save(states, os.path.join(self.args.log, 'checkpoint_{}.pth'.format(step)))
torch.save(states, os.path.join(self.args.log, 'checkpoint.pth'))
def train(self):
if self.config.data.random_flip is False:
tran_transform = test_transform = transforms.Compose([
transforms.Resize(self.config.data.image_size),
transforms.ToTensor()
])
else:
tran_transform = transforms.Compose([
transforms.Resize(self.config.data.image_size),
transforms.RandomHorizontalFlip(p=0.5),
transforms.ToTensor()
])
test_transform = transforms.Compose([
transforms.Resize(self.config.data.image_size),
transforms.ToTensor()
])
if self.config.data.dataset == 'CIFAR10':
dataset = CIFAR10(os.path.join(self.args.run, 'datasets', 'cifar10'), train=True, download=True,
transform=tran_transform)
test_dataset = CIFAR10(os.path.join(self.args.run, 'datasets', 'cifar10_test'), train=False, download=True,
transform=test_transform)
elif self.config.data.dataset == 'MNIST':
dataset = MNIST(os.path.join(self.args.run, 'datasets', 'mnist'), train=True, download=True,
transform=tran_transform)
test_dataset = MNIST(os.path.join(self.args.run, 'datasets', 'mnist_test'), train=False, download=True,
transform=test_transform)
elif self.config.data.dataset == 'CELEBA':
if self.config.data.random_flip:
dataset = CelebA(root=os.path.join(self.args.run, 'datasets', 'celeba'), split='train',
transform=transforms.Compose([
transforms.CenterCrop(140),
transforms.Resize(self.config.data.image_size),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
]), download=True)
else:
dataset = CelebA(root=os.path.join(self.args.run, 'datasets', 'celeba'), split='train',
transform=transforms.Compose([
transforms.CenterCrop(140),
transforms.Resize(self.config.data.image_size),
transforms.ToTensor(),
]), download=True)
test_dataset = CelebA(root=os.path.join(self.args.run, 'datasets', 'celeba_test'), split='test',
transform=transforms.Compose([
transforms.CenterCrop(140),
transforms.Resize(self.config.data.image_size),
transforms.ToTensor(),
]), download=True)
elif self.config.data.dataset == 'SVHN':
dataset = SVHN(os.path.join(self.args.run, 'datasets', 'svhn'), split='train', download=True,
transform=tran_transform)
test_dataset = SVHN(os.path.join(self.args.run, 'datasets', 'svhn_test'), split='test', download=True,
transform=test_transform)
elif self.config.data.dataset == 'NYUv2':
if self.config.data.random_flip is False:
nyu_train_transform = nyu_test_transform = transforms.Compose([
transforms.CenterCrop((400, 400)),
transforms.Resize(32),
transforms.ToTensor()
])
else:
nyu_train_transform = transforms.Compose([
transforms.CenterCrop((400, 400)),
transforms.Resize(32),
transforms.RandomHorizontalFlip(p=0.5),
transforms.ToTensor()
])
nyu_test_transform = transforms.Compose([
transforms.CenterCrop((400, 400)),
transforms.Resize(32),
transforms.ToTensor()
])
dataset = NYUv2(os.path.join(self.args.run, 'datasets', 'nyuv2'), train=True, download=True,
rgb_transform=nyu_train_transform, depth_transform=nyu_train_transform)
test_dataset = NYUv2(os.path.join(self.args.run, 'datasets', 'nyuv2'), train=False, download=True,
rgb_transform=nyu_test_transform, depth_transform=nyu_test_transform)
dataloader = DataLoader(dataset, batch_size=self.config.training.batch_size, shuffle=True,
num_workers=0) # changed num_workers from 4 to 0
test_loader = DataLoader(test_dataset, batch_size=self.config.training.batch_size, shuffle=True,
num_workers=0, drop_last=True) # changed num_workers from 4 to 0
test_iter = iter(test_loader)
self.config.input_dim = self.config.data.image_size ** 2 * self.config.data.channels
tb_path = os.path.join(self.args.run, 'tensorboard', self.args.doc)
if os.path.exists(tb_path):
shutil.rmtree(tb_path)
tb_logger = tensorboardX.SummaryWriter(log_dir=tb_path)
score = CondRefineNetDilated(self.config).to(self.config.device)
score = torch.nn.DataParallel(score)
optimizer = self.get_optimizer(score.parameters())
if self.args.resume_training:
states = torch.load(os.path.join(self.args.log, 'checkpoint.pth'))
score.load_state_dict(states[0])
optimizer.load_state_dict(states[1])
step = 0
sigmas = torch.tensor(
np.exp(np.linspace(np.log(self.config.model.sigma_begin), np.log(self.config.model.sigma_end),
self.config.model.num_classes))).float().to(self.config.device)
for epoch in range(self.config.training.n_epochs):
for i, (X, y) in enumerate(dataloader):
step += 1
score.train()
X = X.to(self.config.device)
X = X / 256. * 255. + torch.rand_like(X) / 256.
if self.config.data.logit_transform:
X = self.logit_transform(X)
if self.config.data.dataset == 'NYUv2':
# concatenate depth map with image
y = y[0]
# code to see resized depth map
# input_gt_depth_image = y[0][0].data.cpu().numpy().astype(np.float32)
# plot.imsave('gt_depth_map_{}.png'.format(i), input_gt_depth_image,
# cmap="viridis")
y = y.to(self.config.device)
X = torch.cat((X, y), 1)
labels = torch.randint(0, len(sigmas), (X.shape[0],), device=X.device)
if self.config.training.algo == 'dsm':
loss = anneal_dsm_score_estimation(score, X, labels, sigmas, self.config.training.anneal_power)
elif self.config.training.algo == 'ssm':
loss = anneal_sliced_score_estimation_vr(score, X, labels, sigmas,
n_particles=self.config.training.n_particles)
optimizer.zero_grad()
loss.backward()
optimizer.step()
tb_logger.add_scalar('loss', loss, global_step=step)
logging.info("step: {}, loss: {}".format(step, loss.item()))
if step >= self.config.training.n_iters:
return 0
if step % 100 == 0:
score.eval()
try:
test_X, test_y = next(test_iter)
except StopIteration:
test_iter = iter(test_loader)
test_X, test_y = next(test_iter)
test_X = test_X.to(self.config.device)
test_X = test_X / 256. * 255. + torch.rand_like(test_X) / 256.
if self.config.data.logit_transform:
test_X = self.logit_transform(test_X)
if self.config.data.dataset == 'NYUv2':
test_y = test_y[0]
test_y = test_y.to(self.config.device)
test_X = torch.cat((test_X, test_y), 1)
test_labels = torch.randint(0, len(sigmas), (test_X.shape[0],), device=test_X.device)
with torch.no_grad():
test_dsm_loss = anneal_dsm_score_estimation(score, test_X, test_labels, sigmas,
self.config.training.anneal_power)
tb_logger.add_scalar('test_dsm_loss', test_dsm_loss, global_step=step)
if step % self.config.training.snapshot_freq == 0:
states = [
score.state_dict(),
optimizer.state_dict(),
]
torch.save(states, os.path.join(self.args.log, 'checkpoint_{}.pth'.format(step)))
torch.save(states, os.path.join(self.args.log, 'checkpoint.pth'))
def train(self):
dataset, test_dataset = get_dataset(self.args, self.config)
dataloader = DataLoader(dataset, batch_size=self.config.training.batch_size, shuffle=True,
num_workers=self.config.data.num_workers)
test_loader = DataLoader(test_dataset, batch_size=self.config.training.batch_size, shuffle=True,
num_workers=self.config.data.num_workers, drop_last=True)
test_iter = iter(test_loader)
self.config.input_dim = self.config.data.image_size ** 2 * self.config.data.channels
tb_logger = self.config.tb_logger
score = get_model(self.config)
score = torch.nn.DataParallel(score)
optimizer = get_optimizer(self.config, score.parameters())
start_epoch = 0
step = 0
if self.config.model.ema:
ema_helper = EMAHelper(mu=self.config.model.ema_rate)
ema_helper.register(score)
if self.args.resume_training:
states = torch.load(os.path.join(self.args.log_path, 'checkpoint.pth'))
score.load_state_dict(states[0])
### Make sure we can resume with different eps
states[1]['param_groups'][0]['eps'] = self.config.optim.eps
optimizer.load_state_dict(states[1])
start_epoch = states[2]
step = states[3]
if self.config.model.ema:
ema_helper.load_state_dict(states[4])
sigmas = get_sigmas(self.config)
if self.config.training.log_all_sigmas:
### Commented out training time logging to save time.
test_loss_per_sigma = [None for _ in range(len(sigmas))]
def hook(loss, labels):
# for i in range(len(sigmas)):
# if torch.any(labels == i):
# test_loss_per_sigma[i] = torch.mean(loss[labels == i])
pass
def tb_hook():
# for i in range(len(sigmas)):
# if test_loss_per_sigma[i] is not None:
# tb_logger.add_scalar('test_loss_sigma_{}'.format(i), test_loss_per_sigma[i],
# global_step=step)
pass
def test_hook(loss, labels):
for i in range(len(sigmas)):
if torch.any(labels == i):
test_loss_per_sigma[i] = torch.mean(loss[labels == i])
def test_tb_hook():
for i in range(len(sigmas)):
if test_loss_per_sigma[i] is not None:
tb_logger.add_scalar('test_loss_sigma_{}'.format(i), test_loss_per_sigma[i],
global_step=step)
else:
hook = test_hook = None
def tb_hook():
pass
def test_tb_hook():
pass
for epoch in range(start_epoch, self.config.training.n_epochs):
for i, (X, y) in enumerate(dataloader):
score.train()
step += 1
X = X.to(self.config.device)
X = data_transform(self.config, X)
loss = anneal_dsm_score_estimation(score, X, sigmas, None,
self.config.training.anneal_power,
hook)
tb_logger.add_scalar('loss', loss, global_step=step)
tb_hook()
logging.info("step: {}, loss: {}".format(step, loss.item()))
optimizer.zero_grad()
loss.backward()
optimizer.step()
if self.config.model.ema:
ema_helper.update(score)
if step >= self.config.training.n_iters:
return 0
if step % 100 == 0:
if self.config.model.ema:
test_score = ema_helper.ema_copy(score)
else:
test_score = score
test_score.eval()
try:
test_X, test_y = next(test_iter)
except StopIteration:
test_iter = iter(test_loader)
test_X, test_y = next(test_iter)
test_X = test_X.to(self.config.device)
test_X = data_transform(self.config, test_X)
with torch.no_grad():
test_dsm_loss = anneal_dsm_score_estimation(test_score, test_X, sigmas, None,
self.config.training.anneal_power,
hook=test_hook)
tb_logger.add_scalar('test_loss', test_dsm_loss, global_step=step)
test_tb_hook()
logging.info("step: {}, test_loss: {}".format(step, test_dsm_loss.item()))
del test_score
if step % self.config.training.snapshot_freq == 0:
states = [
score.state_dict(),
optimizer.state_dict(),
epoch,
step,
]
if self.config.model.ema:
states.append(ema_helper.state_dict())
torch.save(states, os.path.join(self.args.log_path, 'checkpoint_{}.pth'.format(step)))
torch.save(states, os.path.join(self.args.log_path, 'checkpoint.pth'))
if self.config.training.snapshot_sampling:
if self.config.model.ema:
test_score = ema_helper.ema_copy(score)
else:
test_score = score
test_score.eval()
## Different part from NeurIPS 2019.
## Random state will be affected because of sampling during training time.
init_samples = torch.rand(36, self.config.data.channels,
self.config.data.image_size, self.config.data.image_size,
device=self.config.device)
init_samples = data_transform(self.config, init_samples)
all_samples = anneal_Langevin_dynamics(init_samples, test_score, sigmas.cpu().numpy(),
self.config.sampling.n_steps_each,
self.config.sampling.step_lr,
final_only=True, verbose=True,
denoise=self.config.sampling.denoise)
sample = all_samples[-1].view(all_samples[-1].shape[0], self.config.data.channels,
self.config.data.image_size,
self.config.data.image_size)
sample = inverse_data_transform(self.config, sample)
image_grid = make_grid(sample, 6)
save_image(image_grid,
os.path.join(self.args.log_sample_path, 'image_grid_{}.png'.format(step)))
torch.save(sample, os.path.join(self.args.log_sample_path, 'samples_{}.pth'.format(step)))
del test_score
del all_samples
def train(self):
if self.config.data.random_flip is False:
tran_transform = test_transform = transforms.Compose([
transforms.Resize(self.config.data.image_size),
transforms.ToTensor()
])
else:
tran_transform = transforms.Compose([
transforms.Resize(self.config.data.image_size),
transforms.RandomHorizontalFlip(p=0.5),
transforms.ToTensor()
])
test_transform = transforms.Compose([
transforms.Resize(self.config.data.image_size),
transforms.ToTensor()
])
if self.config.data.dataset == 'CIFAR10':
dataset = CIFAR10(os.path.join(self.args.run, 'datasets',
'cifar10'),
train=True,
download=True,
transform=tran_transform)
test_dataset = CIFAR10(os.path.join(self.args.run, 'datasets',
'cifar10_test'),
train=False,
download=True,
transform=test_transform)
elif self.config.data.dataset == 'MNIST':
dataset = MNIST(os.path.join(self.args.run, 'datasets', 'mnist'),
train=True,
download=True,
transform=tran_transform)
test_dataset = MNIST(os.path.join(self.args.run, 'datasets',
'mnist_test'),
train=False,
download=True,
transform=test_transform)
elif self.config.data.dataset == 'CELEBA':
if self.config.data.random_flip:
dataset = CelebA(
root=os.path.join(self.args.run, 'datasets', 'celeba'),
split='train',
transform=transforms.Compose([
transforms.CenterCrop(140),
transforms.Resize(self.config.data.image_size),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
]),
download=False)
else:
dataset = CelebA(
root=os.path.join(self.args.run, 'datasets', 'celeba'),
split='train',
transform=transforms.Compose([
transforms.CenterCrop(140),
transforms.Resize(self.config.data.image_size),
transforms.ToTensor(),
]),
download=False)
test_dataset = CelebA(
root=os.path.join(self.args.run, 'datasets', 'celeba_test'),
split='test',
transform=transforms.Compose([
transforms.CenterCrop(140),
transforms.Resize(self.config.data.image_size),
transforms.ToTensor(),
]),
download=False)
elif self.config.data.dataset == 'SVHN':
dataset = SVHN(os.path.join(self.args.run, 'datasets', 'svhn'),
split='train',
download=True,
transform=tran_transform)
test_dataset = SVHN(os.path.join(self.args.run, 'datasets',
'svhn_test'),
split='test',
download=True,
transform=test_transform)
dataloader = DataLoader(dataset,
batch_size=self.config.training.batch_size,
shuffle=True,
num_workers=4)
test_loader = DataLoader(test_dataset,
batch_size=self.config.training.batch_size,
shuffle=True,
num_workers=4,
drop_last=True)
test_iter = iter(test_loader)
self.config.input_dim = self.config.data.image_size**2 * self.config.data.channels
tb_path = os.path.join(self.args.run, 'tensorboard', self.args.doc)
if os.path.exists(tb_path):
shutil.rmtree(tb_path)
tb_logger = tensorboardX.SummaryWriter(log_dir=tb_path)
score = CondRefineNetDilated(self.config).to(self.config.device)
score = torch.nn.DataParallel(score)
optimizer = self.get_optimizer(score.parameters())
if self.args.resume_training:
states = torch.load(os.path.join(self.args.log, 'checkpoint.pth'))
score.load_state_dict(states[0])
optimizer.load_state_dict(states[1])
step = 0
sigmas = torch.tensor(
np.exp(
np.linspace(np.log(self.config.model.sigma_begin),
np.log(self.config.model.sigma_end),
self.config.model.num_classes))).float().to(
self.config.device)
time_record = []
for epoch in range(self.config.training.n_epochs):
for i, (X, y) in enumerate(dataloader):
step += 1
score.train()
X = X.to(self.config.device)
X = X / 256. * 255. + torch.rand_like(X) / 256.
if self.config.data.logit_transform:
X = self.logit_transform(X)
labels = torch.randint(0,
len(sigmas), (X.shape[0], ),
device=X.device)
if self.config.training.algo == 'dsm':
t = time.time()
loss = anneal_dsm_score_estimation(
score, X, labels, sigmas,
self.config.training.anneal_power)
elif self.config.training.algo == 'dsm_tracetrick':
t = time.time()
loss = anneal_dsm_score_estimation_TraceTrick(
score, X, labels, sigmas,
self.config.training.anneal_power)
elif self.config.training.algo == 'ssm':
t = time.time()
loss = anneal_sliced_score_estimation_vr(
score,
X,
labels,
sigmas,
n_particles=self.config.training.n_particles)
elif self.config.training.algo == 'esm_scorenet':
t = time.time()
loss = anneal_ESM_scorenet(
score,
X,
labels,
sigmas,
n_particles=self.config.training.n_particles)
elif self.config.training.algo == 'esm_scorenet_VR':
t = time.time()
loss = anneal_ESM_scorenet_VR(
score,
X,
labels,
sigmas,
n_particles=self.config.training.n_particles)
optimizer.zero_grad()
loss.backward()
optimizer.step()
t = time.time() - t
time_record.append(t)
if step >= self.config.training.n_iters:
return 0
if step % 100 == 0:
tb_logger.add_scalar('loss', loss, global_step=step)
logging.info(
"step: {}, loss: {}, time per step: {:.3f} +- {:.3f} ms"
.format(step, loss.item(),
np.mean(time_record) * 1e3,
np.std(time_record) * 1e3))
# if step % 2000 == 0:
# score.eval()
# try:
# test_X, test_y = next(test_iter)
# except StopIteration:
# test_iter = iter(test_loader)
# test_X, test_y = next(test_iter)
# test_X = test_X.to(self.config.device)
# test_X = test_X / 256. * 255. + torch.rand_like(test_X) / 256.
# if self.config.data.logit_transform:
# test_X = self.logit_transform(test_X)
# test_labels = torch.randint(0, len(sigmas), (test_X.shape[0],), device=test_X.device)
# #if self.config.training.algo == 'dsm':
# with torch.no_grad():
# test_dsm_loss = anneal_dsm_score_estimation(score, test_X, test_labels, sigmas,
# self.config.training.anneal_power)
# tb_logger.add_scalar('test_dsm_loss', test_dsm_loss, global_step=step)
# logging.info("step: {}, test dsm loss: {}".format(step, test_dsm_loss.item()))
# elif self.config.training.algo == 'ssm':
# test_ssm_loss = anneal_sliced_score_estimation_vr(score, test_X, test_labels, sigmas,
# n_particles=self.config.training.n_particles)
# tb_logger.add_scalar('test_ssm_loss', test_ssm_loss, global_step=step)
# logging.info("step: {}, test ssm loss: {}".format(step, test_ssm_loss.item()))
if step >= 140000 and step % self.config.training.snapshot_freq == 0:
states = [
score.state_dict(),
optimizer.state_dict(),
]
torch.save(
states,
os.path.join(self.args.log,
'checkpoint_{}.pth'.format(step)))
torch.save(states,
os.path.join(self.args.log, 'checkpoint.pth'))
def run(self):
self.config.input_dim = self.config.data.channels * self.config.data.image_size ** 2
dataloader, testloader = self.get_dataloader(bs=self.config.training.batch_size, training=True)
sample_loader = self.get_dataloader(bs=36)
sigmas = self.get_sigmas(training=True)
sampling_sigmas = self.get_sigmas(npy=True)
score = self.get_model()
optimizer = self.get_optimizer(score.parameters())
if self.config.model.ema:
ema_helper = EMAHelper(mu=self.config.model.ema_rate)
ema_helper.register(score)
else:
ema_helper = DummyEMA()
if self.args.adversarial:
D = self._setup_discriminator()
optimizerD = self.get_optimizer(D.parameters(), adv=True)
D_loss_function, G_loss_function = adv_loss(self.config.adversarial.adv_loss, self.config.training, self.args.device)
else:
D = optimizerD = D_loss_function = G_loss_function = None
if self.args.resume_training:
states = torch.load(os.path.join(self.args.log_path, 'checkpoint.pth'))
score.load_state_dict(states[0])
### Make sure we can resume with different eps
states[1]['param_groups'][0]['eps'] = self.config.optim.eps
optimizer.load_state_dict(states[1])
start_epoch = states[2]
step = states[3]
ema_helper.load_state_dict(states[4])
if self.args.adversarial:
D.load_state_dict(states[5])
else:
start_epoch = 0
step = 0
hook = Hook(self.args.tb_logger, len(sigmas)) if self.config.training.log_all_sigmas else None
testlogger = Logger(self.args, self.config, sigmas, testloader, ema_helper)
kwargs = {'sigmas': sampling_sigmas, 'final_only': True, 'nsigma': self.config.sampling.nsigma,
'step_lr': self.config.sampling.step_lr, 'target': self.args.target, 'noise_first': self.config.sampling.noise_first}
# Estimate maximum sigma that we should be using
if self.args.compute_approximate_sigma_max:
with torch.no_grad():
current_max_dist = 0
for i, (X, y) in enumerate(dataloader):
X = X.to(self.args.device)
X = data_transform(self.config.data, X)
X_ = X.view(X.shape[0], -1)
max_dist = torch.cdist(X_, X_).max().item()
if current_max_dist < max_dist:
current_max_dist = max_dist
print(current_max_dist)
print('Final, max eucledian distance: {}'.format(current_max_dist))
return ()
D_step = 0
for epoch in range(start_epoch, self.config.training.n_epochs):
for i, (X, y) in enumerate(dataloader):
score.train()
X = X.to(self.args.device)
X = data_transform(self.config.data, X)
##### Discriminator steps #####
if self.args.adversarial:
D_step += 1
""" GAN Discriminator update (at every 'scorenetwork' update)"""
loss_D = self.update_adversarial_discriminator(X, y, sigmas, score, D, optimizerD, D_loss_function)
##### Score Network step #####
if not self.args.adversarial or D_step >= self.config.adversarial.D_steps: # Only update Score network if Discriminator has done all its steps
D_step = 0 # We reset the discriminator counter
step += 1
""" Score network update """
if hook is not None:
hook.update_step(step)
loss_dae, fake_denoised_X, scores_ = anneal_dsm_score_estimation(self.args, score, X, sigmas, hook=hook)
if self.args.adversarial:
# Tells me how 'real' fake_denoised_X looks. loss_adv high means the discriminator found
# it easy to tell they were fake.
if self.config.adversarial.adv_clamp:
fake_denoised_X_ = fake_denoised_X.clamp(0, 1)
else:
fake_denoised_X_ = fake_denoised_X
y_pred = D(X)
y_pred_fake = D(fake_denoised_X_)
loss_adv = self.config.adversarial.lambda_G_gan * G_loss_function(y_pred, y_pred_fake)
loss = self.config.adversarial.lambda_dae * loss_dae + loss_adv
_losses = [loss_D.item(), loss_dae.item(), loss_adv.item()]
else:
loss = loss_dae
_losses = [loss.item()]
optimizer.zero_grad()
loss.backward()
optimizer.step()
ema_helper.update(score)
testlogger.addentry(step, score, _losses)
_force_end = step == self.config.training.n_iters
if step % self.config.training.snapshot_freq == 0 or _force_end:
states = [score.state_dict(), optimizer.state_dict(), epoch, step, ema_helper.state_dict()]
torch.save(states, os.path.join(self.args.log_path, 'checkpoint_{}.pth'.format(step)))
torch.save(states, os.path.join(self.args.log_path, 'checkpoint.pth'))
if _force_end:
return 1
if step % self.config.training.snapshot_sampling_freq == 0 and self.config.training.snapshot_sampling:
test_score = ema_helper.ema_copy(score, self.args.device)
test_score.eval()
kwargs['scorenet'] = test_score
self.sample(sample_loader, saveimages=True, kwargs=kwargs, gridsize=36, bs=36, ckpt_id=step)