def test_multi(self):
"""Translate images using StarGAN trained on multiple datasets."""
# Load the trained generator.
self.restore_model(self.test_iters)
with torch.no_grad():
for i, (x_real, c_org) in enumerate(self.celeba_loader):
# Prepare input images and target domain labels.
x_real = x_real.to(self.device)
c_celeba_list = self.create_labels(c_org, self.c_dim, 'CelebA', self.selected_attrs)
c_rafd_list = self.create_labels(c_org, self.c2_dim, 'RaFD')
zero_celeba = torch.zeros(x_real.size(0), self.c_dim).to(self.device) # Zero vector for CelebA.
zero_rafd = torch.zeros(x_real.size(0), self.c2_dim).to(self.device) # Zero vector for RaFD.
mask_celeba = self.label2onehot(torch.zeros(x_real.size(0)), 2).to(self.device) # Mask vector: [1, 0].
mask_rafd = self.label2onehot(torch.ones(x_real.size(0)), 2).to(self.device) # Mask vector: [0, 1].
# Translate images.
x_fake_list = [x_real]
for c_celeba in c_celeba_list:
c_trg = torch.cat([c_celeba, zero_rafd, mask_celeba], dim=1)
x_fake_list.append(self.G(x_real, c_trg))
for c_rafd in c_rafd_list:
c_trg = torch.cat([zero_celeba, c_rafd, mask_rafd], dim=1)
x_fake_list.append(self.G(x_real, c_trg))
# Save the translated images.
x_concat = torch.cat(x_fake_list, dim=3)
result_path = os.path.join(self.result_dir, '{}-images.jpg'.format(i+1))
save_image(self.denorm(x_concat.data.cpu()), result_path, nrow=1, padding=0)
print('Saved real and fake images into {}...'.format(result_path))
def test(self):
"""Translate images using StarGAN trained on a single dataset."""
# Load the trained generator.
self.restore_model(self.test_iters)
# Set data loader.
if self.dataset == 'CelebA':
data_loader = self.celeba_loader
elif self.dataset == 'RaFD':
data_loader = self.rafd_loader
with torch.no_grad():
for i, (x_real, c_org) in enumerate(data_loader):
# Prepare input images and target domain labels.
x_real = x_real.to(self.device)
c_trg_list = self.create_labels(c_org, self.c_dim, self.dataset, self.selected_attrs)
# Translate images.
x_fake_list = [x_real]
for c_trg in c_trg_list:
x_fake_list.append(self.G(x_real, c_trg))
# Save the translated images.
x_concat = torch.cat(x_fake_list, dim=3)
result_path = os.path.join(self.result_dir, '{}-images.jpg'.format(i+1))
save_image(self.denorm(x_concat.data.cpu()), result_path, nrow=1, padding=0)
print('Saved real and fake images into {}...'.format(result_path))
def save_img_results(imgs_tcpu, fake_imgs, num_imgs,
count, image_dir, summary_writer):
num = cfg.TRAIN.VIS_COUNT
# The range of real_img (i.e., self.imgs_tcpu[i][0:num])
# is changed to [0, 1] by function vutils.save_image
real_img = imgs_tcpu[-1][0:num]
vutils.save_image(
real_img, '%s/real_samples.png' % (image_dir),
normalize=True)
real_img_set = vutils.make_grid(real_img).numpy()
real_img_set = np.transpose(real_img_set, (1, 2, 0))
real_img_set = real_img_set * 255
real_img_set = real_img_set.astype(np.uint8)
sup_real_img = summary.image('real_img', real_img_set)
summary_writer.add_summary(sup_real_img, count)
for i in range(num_imgs):
fake_img = fake_imgs[i][0:num]
# The range of fake_img.data (i.e., self.fake_imgs[i][0:num])
# is still [-1. 1]...
vutils.save_image(
fake_img.data, '%s/count_%09d_fake_samples%d.png' %
(image_dir, count, i), normalize=True)
fake_img_set = vutils.make_grid(fake_img.data).cpu().numpy()
fake_img_set = np.transpose(fake_img_set, (1, 2, 0))
fake_img_set = (fake_img_set + 1) * 255 / 2
fake_img_set = fake_img_set.astype(np.uint8)
sup_fake_img = summary.image('fake_img%d' % i, fake_img_set)
summary_writer.add_summary(sup_fake_img, count)
summary_writer.flush()
def sampleTrue(dataset, imageSize, dataroot, sampleSize, batchSize, saveFolder, workers=4):
print('sampling real images ...')
saveFolder = saveFolder + '0/'
dataset = make_dataset(dataset, dataroot, imageSize)
dataloader = torch.utils.data.DataLoader(
dataset, shuffle=True, batch_size=batchSize, num_workers=int(workers))
if not os.path.exists(saveFolder):
try:
os.makedirs(saveFolder)
except OSError:
pass
iter = 0
for i, data in enumerate(dataloader, 0):
img, _ = data
for j in range(0, len(img)):
vutils.save_image(img[j].mul(0.5).add(
0.5), saveFolder + giveName(iter) + ".png")
iter += 1
if iter >= sampleSize:
break
if iter >= sampleSize:
break
def plot_rec(x, netEC, netEP, netD):
x_c = x[0]
x_p = x[np.random.randint(1, opt.max_step)]
h_c = netEC(x_c)
h_p = netEP(x_p)
# print('h_c shape: ', h_c.shape)
# print('h p shape: ', h_p.shape)
rec = netD([h_c, h_p])
x_c, x_p, rec = x_c.data, x_p.data, rec.data
fname = '%s/rec/rec_test.png' % (opt.log_dir)
comparison = None
for i in range(len(x_c)):
if comparison is None:
comparison = torch.stack([x_c[i], x_p[i], rec[i]])
else:
new_comparison = torch.stack([x_c[i], x_p[i], rec[i]])
comparison = torch.cat([comparison, new_comparison])
print('comparison: ', comparison.shape)
# row_sz = 5
# nplot = 20
# for i in range(0, nplot - row_sz, row_sz):
# row = [[xc, xp, xr] for xc, xp, xr in zip(x_c[i:i + row_sz], x_p[i:i + row_sz], rec[i:i + row_sz])]
# print('row: ', row)
# to_plot.append(list(itertools.chain(*row)))
# print(len(to_plot[0]))
# utils.save_tensors_image(fname, comparison)
if not os.path.exists(os.path.dirname(fname)):
os.makedirs(os.path.dirname(fname))
save_image(comparison.cpu(), fname, nrow=3)
def test(self):
"""Facial attribute transfer on CelebA or facial expression synthesis on RaFD."""
# Load trained parameters
G_path = os.path.join(self.model_save_path, '{}_G.pth'.format(self.test_model))
self.G.load_state_dict(torch.load(G_path))
self.G.eval()
if self.dataset == 'CelebA':
data_loader = self.celebA_loader
else:
data_loader = self.rafd_loader
for i, (real_x, org_c) in enumerate(data_loader):
real_x = self.to_var(real_x, volatile=True)
if self.dataset == 'CelebA':
target_c_list = self.make_celeb_labels(org_c)
else:
target_c_list = []
for j in range(self.c_dim):
target_c = self.one_hot(torch.ones(real_x.size(0)) * j, self.c_dim)
target_c_list.append(self.to_var(target_c, volatile=True))
# Start translations
fake_image_list = [real_x]
for target_c in target_c_list:
fake_image_list.append(self.G(real_x, target_c))
fake_images = torch.cat(fake_image_list, dim=3)
save_path = os.path.join(self.result_path, '{}_fake.png'.format(i+1))
save_image(self.denorm(fake_images.data), save_path, nrow=1, padding=0)
print('Translated test images and saved into "{}"..!'.format(save_path))
def saver(state):
if state[torchbearer.BATCH] == 0:
data = state[torchbearer.X]
recon_batch = state[torchbearer.Y_PRED]
comparison = torch.cat([data[:num_images],
recon_batch.view(128, 1, 28, 28)[:num_images]])
save_image(comparison.cpu(),
str(folder) + 'reconstruction_' + str(state[torchbearer.EPOCH]) + '.png', nrow=num_images)
def sample_image(n_row, batches_done):
"""Saves a grid of generated digits ranging from 0 to n_classes"""
# Sample noise
z = Variable(FloatTensor(np.random.normal(0, 1, (n_row**2, opt.latent_dim))))
# Get labels ranging from 0 to n_classes for n rows
labels = np.array([num for _ in range(n_row) for num in range(n_row)])
labels = Variable(LongTensor(labels))
gen_imgs = generator(z, labels)
save_image(gen_imgs.data, 'images/%d.png' % batches_done, nrow=n_row, normalize=True)
def _train(self, epoch):
"""Perform the actual train."""
# put model into train mode
self.d_model.train()
# TODO: why?
cp_loader = deepcopy(self.train_loader)
if self.verbose:
progress_bar = tqdm(total=len(cp_loader),
desc='Current Epoch',
file=sys.stdout,
leave=False,
ncols=75,
position=0,
unit=' Batch')
else:
progress_bar = None
real_label = 1
fake_label = 0
for batch_idx, inputs in enumerate(cp_loader):
# Update Discriminator network maximize log(D(x)) + log(1 - D(G(z)))
# train with real
self.optimizer_d.zero_grad()
inputs = inputs.to(self.device)
batch_size = inputs.size(0)
outputs = self.d_model(inputs)
label = torch.full((batch_size,), real_label, device=self.device)
loss_d_real = self.loss_function(outputs, label)
loss_d_real.backward()
# train with fake
noise = torch.randn((batch_size, self.g_model.nz, 1, 1,), device=self.device)
fake_outputs = self.g_model(noise)
label.fill_(fake_label)
outputs = self.d_model(fake_outputs.detach())
loss_g_fake = self.loss_function(outputs, label)
loss_g_fake.backward()
self.optimizer_d.step()
# (2) Update G network: maximize log(D(G(z)))
self.g_model.zero_grad()
label.fill_(real_label)
outputs = self.d_model(fake_outputs)
loss_g = self.loss_function(outputs, label)
loss_g.backward()
self.optimizer_g.step()
if self.verbose:
if batch_idx % 10 == 0:
progress_bar.update(10)
if self.out_f is not None and batch_idx % 100 == 0:
fake = self.g_model(self.sample_noise)
vutils.save_image(
fake.detach(),
'%s/fake_samples_epoch_%03d.png' % (self.out_f, epoch),
normalize=True)
if self.verbose:
progress_bar.close()
def sample_images(batches_done):
"""Saves a generated sample from the test set"""
imgs = next(iter(val_dataloader))
real_A = Variable(imgs['A'].type(Tensor))
fake_B = G_AB(real_A)
real_B = Variable(imgs['B'].type(Tensor))
fake_A = G_BA(real_B)
img_sample = torch.cat((real_A.data, fake_B.data,
real_B.data, fake_A.data), 0)
save_image(img_sample, 'images/%s/%s.png' % (opt.dataset_name, batches_done), nrow=5, normalize=True)
def save_image(img):
post = transforms.Compose([transforms.Lambda(lambda x: x.mul_(1./255)),
transforms.Normalize(mean=[-0.40760392, -0.45795686, -0.48501961], std=[1,1,1]),
transforms.Lambda(lambda x: x[torch.LongTensor([2,1,0])]), #turn to RGB
])
img = post(img)
img = img.clamp_(0,1)
vutils.save_image(img,
'%s/transfer.png' % (opt.outf),
normalize=True)
return
def reconstruction_loss(self, images, input, size_average=True):
# Get the lengths of capsule outputs.
v_mag = torch.sqrt((input**2).sum(dim=2))
# Get index of longest capsule output.
_, v_max_index = v_mag.max(dim=1)
v_max_index = v_max_index.data
# Use just the winning capsule's representation (and zeros for other capsules) to reconstruct input image.
batch_size = input.size(0)
all_masked = [None] * batch_size
for batch_idx in range(batch_size):
# Get one sample from the batch.
input_batch = input[batch_idx]
# Copy only the maximum capsule index from this batch sample.
# This masks out (leaves as zero) the other capsules in this sample.
batch_masked = Variable(torch.zeros(input_batch.size())).cuda()
batch_masked[v_max_index[batch_idx]] = input_batch[v_max_index[batch_idx]]
all_masked[batch_idx] = batch_masked
# Stack masked capsules over the batch dimension.
masked = torch.stack(all_masked, dim=0)
# Reconstruct input image.
masked = masked.view(input.size(0), -1)
output = self.relu(self.reconstruct0(masked))
output = self.relu(self.reconstruct1(output))
output = self.sigmoid(self.reconstruct2(output))
output = output.view(-1, self.image_channels, self.image_height, self.image_width)
# Save reconstructed images occasionally.
if self.reconstructed_image_count % 10 == 0:
if output.size(1) == 2:
# handle two-channel images
zeros = torch.zeros(output.size(0), 1, output.size(2), output.size(3))
output_image = torch.cat([zeros, output.data.cpu()], dim=1)
else:
# assume RGB or grayscale
output_image = output.data.cpu()
vutils.save_image(output_image, "reconstruction.png")
self.reconstructed_image_count += 1
# The reconstruction loss is the sum squared difference between the input image and reconstructed image.
# Multiplied by a small number so it doesn't dominate the margin (class) loss.
error = (output - images).view(output.size(0), -1)
error = error**2
error = torch.sum(error, dim=1) * 0.0005
# Average over batch
if size_average:
error = error.mean()
return error
def save_images(netG, fixed_noise, outputDir, epoch):
'''
Generates a batch of images from the given 'noise'.
Saves 64 of the generated samples to 'outputDir' system path.
Inputs are the network (netG), a 'noise' input, system path to which images will be saved (outputDir) and current 'epoch'.
'''
noise = Variable(fixed_noise)
netG.eval()
fake = netG(noise)
netG.train()
vutils.save_image(
fake.data[0:64, :, :, :], '%s/fake_samples_epoch_%03d.png' % (outputDir, epoch), nrow=8)
def reconstruct_test(self,epoch):
for i,batch in enumerate(self.test_dataloader):
images = batch['image']
images = images.float()
bumps = batch['bump']
bumps = bumps.float()
masks = batch['mask']
masks = masks.float()
images = Variable(images.cuda())
recon_mask, recon = self.Gnet.forward(images)
output = torch.cat((masks,recon_mask.data.cpu(),bumps,recon.data.cpu()),dim=3)
utils.save_image(output, net.outpath + '/'+str(epoch)+'.'+str(i)+'.jpg',nrow=4, normalize=True)
def sample_images(batches_done):
"""Saves a generated sample from the test set"""
imgs = next(iter(val_dataloader))
X1 = Variable(imgs['A'].type(Tensor))
X2 = Variable(imgs['B'].type(Tensor))
_, Z1 = E1(X1)
_, Z2 = E2(X2)
fake_X1 = G1(Z2)
fake_X2 = G2(Z1)
img_sample = torch.cat((X1.data, fake_X2.data,
X2.data, fake_X1.data), 0)
save_image(img_sample, 'images/%s/%s.png' % (opt.dataset_name, batches_done), nrow=5, normalize=True)
def save(self, source, iteration):
save_dir = os.path.join(self.model_dir, "gen_images")
if os.path.exists(save_dir) == False:
os.mkdir(save_dir)
images_file = os.path.join(save_dir, "image_{}.png".format(iteration))
if self.cuda:
source = source.cuda()
source = Variable(source)
outputs = self.gen_model(source)
vutils.save_image(outputs.cpu().data, images_file, normalize=True)
def pp_interp(net, alpha):
"""
Only works with model_resnet_preproc.py as your
architecture!!!
"""
conv2d = net.d.preproc
deconv2d = nn.ConvTranspose2d(16, 3, 3, stride=1, padding=1)
deconv2d = deconv2d.cuda()
deconv2d.weight = conv2d.weight
gz1 = net.sample(bs=128)
gz2 = net.sample(bs=128)
#alpha = net.sample_lambda(gz1.size(0))
gz_mix = alpha*gz1 + (1.-alpha)*gz2
save_image(gz1*0.5 + 0.5, filename="gz1.png")
save_image(gz2*0.5 + 0.5, filename="gz2.png")
save_image(gz_mix*0.5 + 0.5, filename="gz_mix.png")
# Ok, do the mixup in hidden space.
gz1_h = conv2d(gz1)
gz2_h = conv2d(gz2)
#alpha = 0.05
gz_mix_h = alpha*gz1_h + (1.-alpha)*gz2_h
gz_mix_h_dec = deconv2d(gz_mix_h)
save_image(gz_mix_h_dec*0.5 + 0.5, filename="gz_mix_h_dec.png")
print(conv2d.weight == deconv2d.weight)
import pdb
pdb.set_trace()
def closure():
optimizer.zero_grad()
out = resnet(generated)
style_loss = [GramMSELoss().cuda()(out[i],style_target[i])*style_weight[i] for i in range(len(style_target))]
content_loss = nn.MSELoss().cuda()(out[content_layer_num],content_target)
total_loss = 1000 * sum(style_loss) + sum(content_loss)
total_loss.backward()
if iteration[0] % 100 == 0:
print(total_loss)
v_utils.save_image(image_postprocess(generated.data),"./gen_{}.png".format(iteration[0]))
iteration[0] += 1
return total_loss
def generate(self, input_sample=None):
if input_sample is None:
input_sample = torch.randn(self.gen_training_result[1], self.nz, 1, 1, device=self.device)
if not isinstance(input_sample, torch.Tensor) and \
isinstance(input_sample, np.ndarray):
input_sample = torch.from_numpy(input_sample)
if not isinstance(input_sample, torch.Tensor) and \
not isinstance(input_sample, np.ndarray):
raise TypeError("Input should be a torch.tensor or a numpy.ndarray")
self.net_g.eval()
with torch.no_grad():
input_sample = input_sample.to(self.device)
generated_fake = self.net_g(input_sample)
vutils.save_image(generated_fake.detach(),
'%s/evaluation.png' % self.gen_training_result[0],
normalize=True)
def test(self, epoch):
self.set_train(is_train=False)
test_loss = 0
for i, (x, _) in enumerate(self.test_loader):
with torch.no_grad():
recon_x = self.model_eval(x)[0]
test_loss += self.compute_loss_and_gradient(x)
if i == 0:
n = min(x.size(0), 8)
comparison = torch.cat([x[:n],
recon_x.reshape(self.args.batch_size, 1, 28, 28)[:n]])
save_image(comparison.detach().cpu(),
os.path.join(OUTPUT_DIR, 'reconstruction_' + str(epoch) + '.png'),
nrow=n)
test_loss /= len(self.test_loader.dataset)
print('Test set loss: {:.4f}'.format(test_loss))
def test(epoch):
model.eval()
test_loss = 0
with torch.no_grad():
for i, (data, _) in enumerate(test_loader):
data = data.to(device)
recon_batch, mu, logvar = model(data)
test_loss += loss_function(recon_batch, data, mu, logvar).item()
if i == 0:
n = min(data.size(0), 8)
comparison = torch.cat([data[:n],
recon_batch.view(args.batch_size, 1, 28, 28)[:n]])
save_image(comparison.cpu(),
'results/reconstruction_' + str(epoch) + '.png', nrow=n)
test_loss /= len(test_loader.dataset)
print('====> Test set loss: {:.4f}'.format(test_loss))
def test(epoch):
model.eval()
test_loss = 0
for i, (data, _) in enumerate(test_loader):
if args.cuda:
data = data.cuda()
data = Variable(data, volatile=True)
recon_batch, mu, logvar = model(data)
test_loss += loss_function(recon_batch, data, mu, logvar).data[0]
if i == 0:
n = min(data.size(0), 8)
comparison = torch.cat([data[:n],
recon_batch.view(args.batch_size, 1, 28, 28)[:n]])
save_image(comparison.data.cpu(),
'results/reconstruction_' + str(epoch) + '.png', nrow=n)
test_loss /= len(test_loader.dataset)
print('====> Test set loss: {:.4f}'.format(test_loss))
def __call__(self, image):
self.iteration += 1
if self.noiseLevel > 0.0:
a = image.view(-1)
numNoiseBits = int(a.shape[0] * self.noiseLevel)
permutedIndices = np.random.permutation(a.shape[0])
a[permutedIndices[0:numNoiseBits // 2]] = self.highValue
a[permutedIndices[numNoiseBits // 2:numNoiseBits]] = self.lowValue
# Save a subset of the images for debugging
if self.logDir is not None:
if np.random.random() <= self.logProbability:
outfile = os.path.join(self.logDir,
"im_noise_" + str(int(self.noiseLevel * 100)) + "_"
+ str(self.iteration).rjust(6, '0') + ".png")
save_image(image, outfile)
return image
def plot_analogy(x, netEC, netEP, netD):
x_c = x[0]
h_c = netEC(x_c)
nrow = opt.batch_size
row_sz = opt.max_step
to_plot = []
zeros = torch.zeros(opt.channels, opt.image_width, opt.image_width)
# to_plot.append(zeros)
for i in range(nrow):
to_plot.append(x[0][i])
to_plot = torch.stack(to_plot)
for j in range(0, row_sz):
h_p = netEP(x[j])
for i in range(nrow):
h_p[i] = h_p[0]
rec = netD([h_c, Variable(h_p)])
print('rec shape: ', rec.shape)
to_plot = torch.cat([to_plot, rec])
originals = []
for i in range(row_sz):
originals.append(x[i][0])
originals = [x[0][0]] + originals
print('original: ', len(originals))
print('len(to_plot): ', len(to_plot))
plt_list = []
for i in range(len(to_plot)):
if i % nrow == 0:
print('int(i / nrow): ', int(i / nrow))
plt_list.append(originals[int(i / nrow)])
plt_list.append(to_plot[i])
to_plot = torch.stack(plt_list)
fname = '%s/rec/analogy_test.png' % (opt.log_dir)
if not os.path.exists(os.path.dirname(fname)):
os.makedirs(os.path.dirname(fname))
save_image(to_plot, fname, nrow=(nrow+1))
def sample_images(batches_done):
"""Saves a generated sample from the validation set"""
imgs = next(iter(val_dataloader))
img_samples = None
for img_A, img_B in zip(imgs['A'], imgs['B']):
# Repeat input image by number of channels
real_A = img_A.view(1, *img_A.shape).repeat(8, 1, 1, 1)
real_A = Variable(real_A.type(Tensor))
# Get interpolated noise [-1, 1]
sampled_z = np.repeat(np.linspace(-1, 1, 8)[:, np.newaxis], opt.latent_dim, 1)
sampled_z = Variable(Tensor(sampled_z))
# Generator samples
fake_B = generator(real_A, sampled_z)
# Concatenate samples horisontally
fake_B = torch.cat([x for x in fake_B.data.cpu()], -1)
img_sample = torch.cat((img_A, fake_B), -1)
img_sample = img_sample.view(1, *img_sample.shape)
# Cocatenate with previous samples vertically
img_samples = img_sample if img_samples is None else torch.cat((img_samples, img_sample), -2)
save_image(img_samples, 'images/%s/%s.png' % (opt.dataset_name, batches_done), nrow=5, normalize=True)
def sample_images(batches_done):
"""Saves a generated sample from the validation set"""
imgs = next(iter(val_dataloader))
img_samples = None
for img1, img2 in zip(imgs['A'], imgs['B']):
# Create copies of image
X1 = img1.unsqueeze(0).repeat(opt.style_dim, 1, 1, 1)
X1 = Variable(X1.type(Tensor))
# Get interpolated style codes
s_code = np.repeat(np.linspace(-1, 1, opt.style_dim)[:, np.newaxis], opt.style_dim, 1)
s_code = Variable(Tensor(s_code))
# Generate samples
c_code_1, _ = Enc1(X1)
X12 = Dec2(c_code_1, s_code)
# Concatenate samples horisontally
X12 = torch.cat([x for x in X12.data.cpu()], -1)
img_sample = torch.cat((img1, X12), -1).unsqueeze(0)
# Concatenate with previous samples vertically
img_samples = img_sample if img_samples is None else torch.cat((img_samples, img_sample), -2)
save_image(img_samples, 'images/%s/%s.png' % (opt.dataset_name, batches_done), nrow=5, normalize=True)
def test(self):
self.load_checkpoint()
self.G.to(self.device)
tqdm_loader = tqdm(self.data_loader.test_loader, total=self.data_loader.test_iterations,
desc='Testing at checkpoint {}'.format(self.config.checkpoint))
self.G.eval()
with torch.no_grad():
for i, (x_real, c_org) in enumerate(tqdm_loader):
x_real = x_real.to(self.device)
c_trg_list = self.create_labels(c_org, self.config.attrs)
x_fake_list = [x_real]
for c_trg in c_trg_list:
attr_diff = c_trg - c_org
x_fake_list.append(self.G(x_real, attr_diff.to(self.device)))
x_concat = torch.cat(x_fake_list, dim=3)
result_path = os.path.join(self.config.result_dir, 'sample_{}.jpg'.format(i + 1))
save_image(self.denorm(x_concat.data.cpu()), result_path, nrow=1, padding=0)
def stylization(p_wct, content_image_path, style_image_path, content_seg_path, style_seg_path, output_image_path,
cuda):
# Load image
cont_img = Image.open(content_image_path).convert('RGB')
styl_img = Image.open(style_image_path).convert('RGB')
try:
cont_seg = Image.open(content_seg_path)
styl_seg = Image.open(style_seg_path)
except:
cont_seg = []
styl_seg = []
cont_img = transforms.ToTensor()(cont_img).unsqueeze(0)
styl_img = transforms.ToTensor()(styl_img).unsqueeze(0)
if cuda:
cont_img = cont_img.cuda(0)
styl_img = styl_img.cuda(0)
p_wct.cuda(0)
cont_img = Variable(cont_img, volatile=True)
styl_img = Variable(styl_img, volatile=True)
cont_seg = np.asarray(cont_seg)
styl_seg = np.asarray(styl_seg)
with Timer("Elapsed time in stylization: %f"):
stylized_img = p_wct.transform(cont_img, styl_img, cont_seg, styl_seg)
utils.save_image(stylized_img.data.cpu().float(), output_image_path, nrow=1)
with Timer("Elapsed time in propagation: %f"):
out_img = p_pro.process(output_image_path, content_image_path)
out_img.save(output_image_path)
if not cuda:
print("NotImplemented: The CPU version of smooth filter has not been implemented currently.")
return
with Timer("Elapsed time in post processing: %f"):
out_img = smooth_filter(output_image_path, content_image_path, f_radius=15, f_edge=1e-1)
out_img.save(output_image_path)
def sampleFake(netG, nz, sampleSize, batchSize, saveFolder):
print('sampling fake images ...')
saveFolder = saveFolder + '0/'
try:
os.makedirs(saveFolder)
except OSError:
pass
noise = torch.FloatTensor(batchSize, nz, 1, 1).cuda()
iter = 0
for i in range(0, 1 + sampleSize // batchSize):
noise.data.normal_(0, 1)
fake = netG(noise)
for j in range(0, len(fake.data)):
if iter < sampleSize:
vutils.save_image(fake.data[j].mul(0.5).add(
0.5), saveFolder + giveName(iter) + ".png")
iter += 1
if iter >= sampleSize:
break
def test(self):
"""
:return:
"""
self.G.eval()
for i, (input, target) in enumerate(self.dataloader):
input_A = Variable(input.cuda())
input_B = Variable(target.cuda())
# make downsampled input
self.real_imgs_A = [input_A]
for _ in range(self.num_D - 1):
real_img = self.downsample(self.real_imgs_A[-1])
self.real_imgs_A.append(real_img)
self.real_imgs_A.reverse()
self.real_imgs_B = [input_B]
for _ in range(self.num_D - 1):
real_img = self.downsample(self.real_imgs_B[-1])
self.real_imgs_B.append(real_img)
self.real_imgs_B.reverse()
# (1) generate fake images
self.fake_imgs = self.G(input_A)
vutils.save_image(input_A.data,
'%s/real_samples_A_%03d.png' % (self.config['outf_test'], i),
normalize=True, padding=0)
vutils.save_image(input_B.data,
'%s/real_samples_B_%03d.png' % (self.config['outf_test'], i),
normalize=True, padding=0)
vutils.save_image(self.fake_imgs[0].data,
'%s/fake_samples_0_%3d.png' % (self.config['outf_test'], i),
normalize=True, padding=0)
vutils.save_image(self.fake_imgs[1].data,
'%s/fake_samples_1_%03d.png' % (self.config['outf_test'], i),
normalize=True, padding=0)
vutils.save_image(self.fake_imgs[2].data,
'%s/fake_samples_2_%03d.png' % (self.config['outf_test'], i),
normalize=True, padding=0)
# For KL divergence, see Appendix B in VAE paper or http://yunjey47.tistory.com/43
reconst_loss = F.binary_cross_entropy(x_reconst, x, size_average=False)
kl_div = -0.5 * torch.sum(1 + log_var - mu.pow(2) - log_var.exp())
# Backprop and optimize
loss = reconst_loss + kl_div
optimizer.zero_grad()
loss.backward()
optimizer.step()
if (i + 1) % 10 == 0:
print(
"Epoch[{}/{}], Step [{}/{}], Reconst Loss: {:.4f}, KL Div: {:.4f}"
.format(epoch + 1, num_epochs, i + 1, len(data_loader),
reconst_loss.item(), kl_div.item()))
with torch.no_grad():
# Save the sampled images
z = torch.randn(batch_size, z_dim).to(device)
out = model.decode(z).view(-1, 1, 28, 28)
save_image(
out, os.path.join(sample_dir, 'sampled-{}.png'.format(epoch + 1)))
# Save the reconstructed images
out, _, _ = model(x)
x_concat = torch.cat([x.view(-1, 1, 28, 28),
out.view(-1, 1, 28, 28)],
dim=3)
save_image(
x_concat,
os.path.join(sample_dir, 'reconst-{}.png'.format(epoch + 1)))
# ----------------------
# Train Discriminators
# ----------------------
optimizer_D.zero_grad()
# Determine validity of real and generated images
validity1_real, validity2_real = coupled_discriminators(imgs1, imgs2)
validity1_fake, validity2_fake = coupled_discriminators(gen_imgs1.detach(), gen_imgs2.detach())
d_loss = (
adversarial_loss(validity1_real, valid)
+ adversarial_loss(validity1_fake, fake)
+ adversarial_loss(validity2_real, valid)
+ adversarial_loss(validity2_fake, fake)
) / 4
d_loss.backward()
optimizer_D.step()
print(
"[Epoch %d/%d] [Batch %d/%d] [D loss: %f] [G loss: %f]"
% (epoch, opt.n_epochs, i, len(dataloader1), d_loss.item(), g_loss.item())
)
batches_done = epoch * len(dataloader1) + i
if batches_done % opt.sample_interval == 0:
gen_imgs = torch.cat((gen_imgs1.data, gen_imgs2.data), 0)
save_image(gen_imgs, "images/%d.png" % batches_done, nrow=8, normalize=True)
def test(self):
""" Test GANomaly model.
Args:
dataloader ([type]): Dataloader for the test set
Raises:
IOError: Model weights not found.
"""
with torch.no_grad():
# Load the weights of netg and netd.
if self.opt.load_weights:
path = "./output/{}/{}/train/weights/netG.pth".format(self.name.lower(), self.opt.dataset)
pretrained_dict = torch.load(path)['state_dict']
try:
self.netg.load_state_dict(pretrained_dict)
except IOError:
raise IOError("netG weights not found")
print(' Loaded weights.')
self.opt.phase = 'test'
#self.opt.showProcess.setValue(80)
# Create big error tensor for the test set.
self.an_scores = torch.zeros(size=(len(self.dataloader['test'].dataset),), dtype=torch.float32, device=self.device)
self.gt_labels = torch.zeros(size=(len(self.dataloader['test'].dataset),), dtype=torch.long, device=self.device)
self.latent_i = torch.zeros(size=(len(self.dataloader['test'].dataset), self.opt.nz), dtype=torch.float32, device=self.device)
self.latent_o = torch.zeros(size=(len(self.dataloader['test'].dataset), self.opt.nz), dtype=torch.float32, device=self.device)
# print(" Testing model %s." % self.name)
self.times = []
self.total_steps = 0
epoch_iter = 0
for i, data in enumerate(self.dataloader['test'], 0):
self.total_steps += self.opt.batchsize
epoch_iter += self.opt.batchsize
time_i = time.time()
self.set_input(data)
self.fake, latent_i, latent_o = self.netg(self.input)
error = torch.mean(torch.pow((latent_i-latent_o), 2), dim=1)
time_o = time.time()
self.an_scores[i*self.opt.batchsize : i*self.opt.batchsize+error.size(0)] = error.reshape(error.size(0))
self.gt_labels[i*self.opt.batchsize : i*self.opt.batchsize+error.size(0)] = self.gt.reshape(error.size(0))
self.latent_i [i*self.opt.batchsize : i*self.opt.batchsize+error.size(0), :] = latent_i.reshape(error.size(0), self.opt.nz)
self.latent_o [i*self.opt.batchsize : i*self.opt.batchsize+error.size(0), :] = latent_o.reshape(error.size(0), self.opt.nz)
self.times.append(time_o - time_i)
# Save test images.
if self.opt.save_test_images:
dst = os.path.join(self.opt.outf, self.opt.name, 'test', 'images')
if not os.path.isdir(dst):
os.makedirs(dst)
real, fake, _ = self.get_current_images()
vutils.save_image(real, '%s/real_%03d.eps' % (dst, i+1), normalize=True)
vutils.save_image(fake, '%s/fake_%03d.eps' % (dst, i+1), normalize=True)
# Measure inference time.
self.times = np.array(self.times)
self.times = np.mean(self.times[:100] * 1000)
# Scale error vector between [0, 1]
print(torch.min(self.an_scores))
print(torch.max(self.an_scores))
maxNUM = torch.max(self.an_scores)
minNUM = torch.min(self.an_scores)
self.an_scores = (self.an_scores - torch.min(self.an_scores)) / (torch.max(self.an_scores) - torch.min(self.an_scores))
# auc, eer = roc(self.gt_labels, self.an_scores)
# -------------- 处理阈值 ------------------
print('-------------- 处理阈值 ------------------')
print(len(self.gt_labels))
plt.ion()
scores = {}
##plt.ion()
# Create data frame for scores and labels.
scores['scores'] = self.an_scores
scores['labels'] = self.gt_labels
hist = pd.DataFrame.from_dict(scores)
#hist.to_csv("histogram.csv")
# Filter normal and abnormal scores.
abn_scr = hist.loc[hist.labels == 1]['scores']
nrm_scr = hist.loc[hist.labels == 0]['scores']
# Create figure and plot the distribution.
##fig, axes = plt.subplots(figsize=(4, 4))
b = []
c = []
# for i in range(1000):
# b.append(nrm_scr[i])
# for j in range(1000, 3011):
# c.append(abn_scr[j])
print('asasddda')
print(len(nrm_scr))
print(len(abn_scr))
for i in nrm_scr:
b.append(i)
for j in abn_scr:
c.append(j)
##sns.distplot(nrm_scr, label=r'Normal Scores', color='r', bins=100, hist=True)
##sns.distplot(abn_scr, label=r'Abnormal Scores', color='b', bins=100, hist=True)
nrm = np.zeros((50), dtype=np.int)
minfix = 0.4
abn = np.zeros((50), dtype=np.int)
abmin = 30
for k in np.arange(0, 1, 0.02):
kint = int(k * 50)
for j in range(len(nrm_scr)):
if b[j] >= k and b[j] < (k + 0.02):
nrm[kint] = nrm[kint] + 1
for j in range(len(abn_scr)):
if c[j] >= k and c[j] < (k + 0.02):
abn[kint] = abn[kint] + 1
print(nrm)
print(abn)
# startInd = 3
# for k in range(0,20):
# if abs(nrm[k] - abn[k]) <= 3:
# continue
# else:
# startInd = k
# max_dist = (len(nrm) + len(abn))*0.28
# for k in range(startInd, 20):
# if abs(nrm[k] - abn[k]) < 5:
# #max_dist = abs(nrm[k] - abn[k])
# minfix = round((k / 20) + 0.02, 3)
# break
# for k in range(3, 17):
# # print(nrm[k])
# # print(abn[k])
# # print('----')
# if abs(nrm[k]-abn[k]) > abmin and not (nrm[k] == 0 and abn[k] == 0):
# abmin = abs(nrm[k] - abn[k])
# minfix = round((k / 20) + 0.02, 3)
max_dist = (len(nrm) + len(abn)) * 0.25
for k in range(0,50):
num1 = np.sum(nrm[0:k])
num2 = np.sum(abn[k::])
if (num1 + num2) >= max_dist:
minfix = round((k / 50) + 0.05, 3)
max_dist = num1+num2
proline = minfix
print(proline)
print(self.gt_labels[0:20])
print(self.an_scores[0:20])
print('------------- 处理阈值 END --------------')
# ------------- 处理阈值 END --------------
auc = evaluate(self.gt_labels, self.an_scores, metric=self.opt.metric)
performance = OrderedDict([('Avg Run Time (ms/batch)', self.times), ('AUC', auc)])
if self.opt.display_id > 0 and self.opt.phase == 'test':
counter_ratio = float(epoch_iter) / len(self.dataloader['test'].dataset)
self.visualizer.plot_performance(self.epoch, counter_ratio, performance)
# --- 写入文件 ---
dict_info = {}
dict_info['minVal'] = float(minNUM.item())
dict_info['maxVal'] = float(maxNUM.item())
dict_info['proline'] = float(proline)
dict_info['auc'] = float(auc)
dict_info['Avg Run Time (ms/batch)'] = float(self.times)
#self.opt.showText.append(str(performance));
#self.opt.showProcess.setValue(100)
return performance, dict_info
x_reconst, mu, log_var = model(x)
# Compute reconstruction loss and kl divergence
# For KL divergence, see Appendix B in VAE paper or http://yunjey47.tistory.com/43
reconst_loss = F.binary_cross_entropy(x_reconst, x, reduction='sum')
kl_div = -0.5 * torch.sum(1 + log_var - mu.pow(2) - log_var.exp())
loss = reconst_loss + kl_div
optimizer.zero_grad()
loss.backward()
optimizer.step()
total_reconst_loss += reconst_loss.item()
total_kl_loss += kl_div.item()
total_loss += loss.item()
print(
"Epoch[{}/{}]], Total Loss: {:.4f}, Reconst Loss: {:.4f}, KL Div: {:.4f}"
.format(epoch + 1, num_epochs, total_loss / len(mnist),
total_reconst_loss / len(mnist), total_kl_loss / len(mnist)))
with torch.no_grad():
# Save the sampled images
z = torch.randn(batch_size, z_dim).to(device)
out = model.decode(z).view(-1, 1, 28, 28)
save_image(
out, os.path.join(sample_dir, 'sampled-{}.png'.format(epoch + 1)))
# Save the reconstructed images
save_image(
x_reconst.view(-1, 1, 28, 28),
os.path.join(sample_dir, 'reconst-{}.png'.format(epoch + 1)))
import torch
from utils1 import *
from torchvision.utils import save_image
import torch.nn as nn
from sagan_models import Generator_SA, Discriminator_SA
from igan_models import Generator_INV, Discriminator_INV
from dcgan_models import Generator_DC, Discriminator_DC
from gan_models import Generator_MLP, Discriminator_MLP
path = 'models/mnist_sagan_5.5/'
listdir = os.listdir(path)
for i, pth in enumerate(listdir):
if "_G" in pth:
epoch = pth.split("_G")[0]
print(i)
model = torch.load(os.path.join(path, pth))
# print(model)
torch.manual_seed(0)
# print(torch.randn(1,128))
if not os.path.exists(os.path.join(path, 'img/' + epoch)):
os.makedirs(os.path.join(path, 'img/' + epoch))
for i in range(1000):
fixed_z = tensor2var(torch.randn(1, 128))
fake_images, _, _ = model(fixed_z)
fake_images_new = fake_images[0].view(1, 3, 64, 64)
save_image(
denorm(fake_images_new.data),
os.path.join(path, 'img/' + epoch + '/' + str(i) + '.jpg'))
d_loss.backward()
d_opt.step()
z = tc.randn(mini_batch, latent_sz).view(-1, latent_sz, 1,
1).to(device)
fake_images = G(z)
D_result = D(fake_images).squeeze()
g_loss = loss_func(D_result, real_label)
G.zero_grad()
g_loss.backward()
g_opt.step()
if step % 200 == 0:
print(
'epoch {}/{}, step {}, d_loss {:.4f}, g_loss {:.4f}, Real_score {:.2f}, Fake_score {:.2f}'
.format(ep, epoch_sz, step, d_loss.item(), g_loss.item(),
D(images).mean().item(),
D(fake_images).mean().item()))
if ep + 1 == 1:
out = images
out = image_range(out)
save_image(out, os.path.join(result_path, 'real_img.png'))
out = fake_images
out = image_range(out)
save_image(out, os.path.join(result_path, 'fake_img {}.png'.format(ep)))
tc.save(G.state_dict(), 'G.ckpt')
tc.save(D.state_dict(), 'D.ckpt')
def train(self):
print(f'''Starting training:
Epochs: {self.epochs}
Batch size: {self.bs}
Learning rate: {self.lr}
Checkpoints: {self.save_cp}
Device: {self.device.type}
''')
if self.save_cp:
try:
if not os.path.exists(os.path.join(self.outdir, 'ckpt')):
os.makedirs(os.path.join(self.outdir, 'ckpt'))
print('Created checkpoint directory')
except OSError:
pass
optimizerG = optim.Adam(self.netG.parameters(),
lr=self.lr,
betas=(0.9, 0.999))
optimizerD = optim.Adam(self.netD.parameters(),
lr=self.lr,
betas=(0.9, 0.999))
criterion = nn.BCELoss()
criterionL1 = torch.nn.L1Loss().to(self.device)
criterionL2 = torch.nn.MSELoss().to(self.device)
criterionGAN = GANLoss('vanilla').to(self.device)
# criterionGAN = nn.BCEWithLogitsLoss().to(self.device)
# torch.autograd.set_detect_anomaly(True)
for epoch in range(self.start_epoch, self.epochs):
print(f'D is trained {max(1, 5-epoch)} times')
for i, data in enumerate(self.dataloader):
for _ in range(1): #max(1, 5-epoch)
self.netD.zero_grad()
self.set_requires_grad(self.netD, True)
############################
# Loss_D: L_D = -(log(D(y) + log(1 - D(G(x))))
###########################
# forward
y = data[0].to(self.device) # 320 x 180
x = self.gen_x(y).to(self.device) # 320 x 180
Gx = self.netG(x) # 320 x 180
# train with fake
fake_pair = torch.cat((x, Gx), 1)
D_Gx = self.netD(fake_pair.detach())
errD_fake = criterionGAN(D_Gx, False)
# train with real
real_pair = torch.cat((x, y), 1)
D_y = self.netD(real_pair)
errD_real = criterionGAN(D_y, True)
# backprop
errD = (errD_real + errD_fake) * 0.5
errD.backward()
optimizerD.step()
self.netG.zero_grad()
self.set_requires_grad(
self.netD,
False) # D requires no gradients when optimizing G
############################
# Loss_G_GAN: L_G = -log(D(G(x)) # Fake the D
###########################
D_Gx = self.netD(fake_pair)
errG_GAN = criterionGAN(D_Gx, True)
############################
# Loss_G_C: L_C = ||y - G(x)||_1
###########################
errG_C = criterionL2(Gx, y) * 50
# errG_CM = criterionL2((1-self.mask)*Gx, (1-self.mask)*y) * 50
# backprop
errG = errG_GAN + errG_C
errG.backward()
optimizerG.step()
if i % 100 == 0:
print(
f'[{epoch}/{self.epochs}][{i}/{len(self.dataloader)}] LossD_real: {errD_real.item():.4f} LossD_fake: {errD_fake.item():.4f} LossG_GAN: {errG_GAN.item():.4f} LossG_C: {errG_C.item():.4f}'
)
# Log
print(
f'[{epoch}/{self.epochs}][{i}/{len(self.dataloader)}] LossD: {errD.item():.4f} LossG: {errG.item():.4f}'
)
if self.save_cp:
torch.save(self.netG.state_dict(),
f'{self.outdir}/ckpt/G_epoch{epoch}.pth')
torch.save(self.netD.state_dict(),
f'{self.outdir}/ckpt/D_epoch{epoch}.pth')
vutils.save_image(Gx.detach(),
'%s/impainted_samples_epoch_%03d.png' %
(self.outdir, epoch),
normalize=True)
# model = load_local_or_remote_file(model_path)
state_dict = torch.load(model_path)
model.load_state_dict(state_dict)
model.to(ptu.device)
model.eval()
for color in ["grey"]:
reference_path = "demos/door_demos_v3/demo_v3_%s_0.pkl"%color
traj = np.load("demos/door_demos_v3/demo_v3_%s_0.pkl"%color, allow_pickle=True)[0]
goal_image_flat = traj["observations"][-1]["image_observation"]
goal_image = goal_image_flat.reshape(1, 3, 500, 300).transpose([0, 1, 3, 2]) / 255.0
# goal_image = goal_image[:, ::-1, :, :].copy() # flip bgr
goal_image = goal_image[:, :, 60:300, 30:470]
goal_image_pt = ptu.from_numpy(goal_image)
save_image(goal_image_pt.data.cpu(), 'goal.png', nrow=1)
goal_latent = model.encode(goal_image_pt).detach().cpu().numpy().flatten()
initial_image_flat = traj["observations"][0]["image_observation"]
initial_image = initial_image_flat.reshape(1, 3, 500, 300).transpose([0, 1, 3, 2]) / 255.0
# initial_image = initial_image[:, ::-1, :, :].copy() # flip bgr
initial_image = initial_image[:, :, 60:300, 30:470]
initial_image_pt = ptu.from_numpy(initial_image)
save_image(initial_image_pt.data.cpu(), 'initial.png', nrow=1)
initial_latent = model.encode(initial_image_pt).detach().cpu().numpy().flatten()
print("Finished initial_latent")
reward_params = dict(
goal_latent=goal_latent,
initial_latent=initial_latent,
goal_image=goal_image_flat,
initial_image=initial_image_flat,
def train():
dir_name, save_path = getSavePath()
netG, netD = networks.getGD_SN(args.structure, args.dataset, args.image_size, args.num_features,
dim_z=args.input_dim, bottleneck=args.bottleneck)
if args.ema_trick:
ema_netG_9999 = copy.deepcopy(netG)
if args.reload > 0:
netG.load_state_dict(torch.load(save_path + 'G_epoch{}.pth'.format(args.reload)))
netD.load_state_dict(torch.load(save_path + 'D_epoch{}.pth'.format(args.reload)))
if args.ema_trick:
ema_netG_9999.load_state_dict(
torch.load(save_path + 'emaG0.9999_epoch{}.pth'.format(args.reload), map_location=torch.device('cpu')))
netG.cuda()
netD.cuda()
g_optimizer = torch.optim.Adam(netG.parameters(), lr=args.g_lr, betas=(args.beta1, args.beta2))
d_optimizer = torch.optim.Adam(netD.parameters(), lr=args.d_lr, betas=(args.beta1, args.beta2))
g_losses, d_losses = [], []
grad_normD, grad_normG = [], []
loader = datasets.getDataLoader(args.dataset, args.image_size, batch_size=args.batch_size)
data_iter = iter(loader)
for i in range(1, args.num_iters+1):
if i >= args.lr_decay_start:
utils.decay_lr(g_optimizer, args.num_iters, args.lr_decay_start, args.g_lr)
utils.decay_lr(d_optimizer, args.num_iters, args.lr_decay_start, args.d_lr)
if i <= args.reload:
continue
if i == 1:
torch.save(netG.state_dict(), save_path + 'G_epoch0.pth')
torch.save(netD.state_dict(), save_path + 'D_epoch0.pth')
# G-step
for _ in range(args.g_freq):
try:
x = next(data_iter)[0].cuda().float()
except StopIteration:
data_iter = iter(loader)
x = next(data_iter)[0].cuda().float()
z = torch.randn(args.batch_size, args.input_dim, device=device)
g_optimizer.zero_grad()
x_hat = netG(z)
y_hat = netD(x_hat)
y = netD(x)
g_loss = get_gloss(args.losstype, y_hat, y)
g_losses.append(g_loss.item())
g_loss.backward()
g_optimizer.step()
grad_normG.append(utils.getGradNorm(netG))
if args.ema_trick:
utils.soft_copy_param(ema_netG_9999, netG, 0.9999)
for _ in range(args.d_freq):
try:
x = next(data_iter)[0].cuda().float()
except StopIteration:
data_iter = iter(loader)
x = next(data_iter)[0].cuda().float()
z = torch.randn(args.batch_size, args.input_dim, device=device)
d_optimizer.zero_grad()
x_hat = netG(z).detach()
y_hat = netD(x_hat)
y = netD(x)
d_loss = get_dloss(args.losstype, y_hat, y)
d_losses.append(d_loss.item())
d_loss.backward()
d_optimizer.step()
grad_normD.append(utils.getGradNorm(netD))
netD.proj.weight.data = F.normalize(netD.proj.weight.data, dim=1)
if i % args.print_freq == 0:
print('Iteration: {}; G-Loss: {}; D-Loss: {};'.format(i, g_loss, d_loss))
if i == 1:
save_image((x / 2. + 0.5), os.path.join(dir_name, 'real.pdf'))
if i == 1 or i % args.plot_freq == 0:
plot_x = netG(torch.randn(args.batch_size, args.input_dim, device=device)).data
plot_x = plot_x / 2. + 0.5
save_image(plot_x, os.path.join(dir_name, 'fake_images-{}.pdf'.format(i + 1)))
utils.plot_losses(g_losses, d_losses, grad_normG, grad_normD, dir_name)
if i % args.save_freq == 0:
torch.save(netG.state_dict(), save_path + 'G_epoch{}.pth'.format(i))
torch.save(netD.state_dict(), save_path + 'D_epoch{}.pth'.format(i))
if args.ema_trick:
torch.save(ema_netG_9999.state_dict(), save_path + 'emaG0.9999_epoch{}.pth'.format(i))
def save_images(self, images, shape, filename):
# Additional dimensions are only passed to the shape instance when > 1
dimensions = 1 if len(shape) == 3 else shape[3]
img_view = images.view(images.size(0), dimensions, shape[1], shape[2])
save_image(denorm(img_view.data), filename)
d_B_optimizer.zero_grad()
loss_B_real = criterion_gan(real_B, valid(real_B))
loss_B_fake = criterion_gan(fake_B.detach(), fake(fake_B))
loss_discriminator_B = (loss_B_real + loss_B_fake) / 2
loss_discriminator_B = Variable(loss_discriminator_B,
requires_grad=True)
loss_discriminator_B.backward()
d_B_optimizer.step()
loss_d = (loss_discriminator_A + loss_discriminator_B) / 2
loss = loss_g + loss_d
if i % 50 == 0:
print(
'Epoch [{}/{}], Step [{}/{}], Loss:{:.4f}, G-Loss: {:.4f} , G-R-Loss: {:.4f} , G-I-Loss: {:.4f} , G-Adv-Loss: {:.4f}, D-Loss: {:.4f}, D-A-Loss: {:.4f}, D-B-Loss: {:.4f} '
.format(iteration + 1, args.epoch, i, int(len(train_loader)),
loss.item(), loss_g, loss_recon, loss_identity,
loss_gen, loss_d, loss_discriminator_A,
loss_discriminator_B))
if i % 20 == 0:
plot_list = []
for image in [real_A[0], real_B[0], fake_A[0], fake_B[0]]:
image = image.detach().numpy()
plot_list.append(image)
save_image(torch.tensor(plot_list),
'generated-images/' + str(iteration + 1) + '-' +
str(i) + '.png',
nrow=4)
def train(self):
dis_criterion = nn.BCELoss()
aux_criterion = nn.NLLLoss() # add class loss
input = torch.FloatTensor(self.batch_size, 3, self.image_size,
self.image_size)
noise = torch.FloatTensor(self.batch_size, self.nz, 1, 1)
fixed_noise = torch.FloatTensor(self.batch_size, self.nz, 1,
1).normal_(0, 1)
dis_label = torch.FloatTensor(self.batch_size)
aux_label = torch.LongTensor(self.batch_size) # add class label
real_label = 1
fake_label = 0
if self.cuda:
dis_criterion.cuda()
aux_criterion.cuda()
input, dis_label, aux_label = input.cuda(), dis_label.cuda(
), aux_label.cuda() # add class label
noise, fixed_noise = noise.cuda(), fixed_noise.cuda()
inputv = Variable(input)
noisev = Variable(noise)
fixed_noisev = Variable(fixed_noise)
dis_labelv = Variable(dis_label)
aux_labelv = Variable(aux_label)
# setup optimizer
optimizerD = optim.Adam(self.netD.parameters(),
lr=self.lr,
betas=(self.beta1, 0.999))
optimizerG = optim.Adam(self.netG.parameters(),
lr=self.lr,
betas=(self.beta1, 0.999))
for epoch in range(self.niter):
for i, data in enumerate(self.data_loader, 0):
############################
# (1) Update D network: maximize log(D(x)) + log(1 - D(G(z)))
###########################
for p in self.netD.parameters(): # reset requires_grad
p.requires_grad = True # they are set to False below in netG update
# train with real
self.netD.zero_grad()
real_cpu, c_label = data # add c_label
batch_size = real_cpu.size(0)
if self.cuda:
real_cpu = real_cpu.cuda()
c_label = c_label.cuda() # add c label
inputv.data.resize_as_(real_cpu).copy_(real_cpu)
dis_labelv.data.resize_(batch_size).fill_(real_label)
aux_labelv.data.resize_(batch_size).copy_(c_label)
dis_out, aux_out = self.netD(inputv)
dis_errD_real = dis_criterion(dis_out, dis_labelv)
aux_errD_real = aux_criterion(aux_out, aux_labelv)
errD_real = dis_errD_real + aux_errD_real
errD_real.backward()
D_x = dis_out.data.mean()
# train with fake
noisev.data.resize_(batch_size, self.nz, 1, 1).normal_(0, 1)
c_label = np.random.randint(0, self.nl, batch_size)
noisev_ = np.random.normal(0, 1, (batch_size, self.nz))
class_onehot = np.zeros((batch_size, self.nl))
class_onehot[np.arange(batch_size), c_label] = 1
noisev_[np.arange(batch_size), :self.nl] = class_onehot[
np.arange(batch_size)]
noisev_ = (torch.from_numpy(noisev_))
noisev.data.copy_(noisev_.view(batch_size, self.nz, 1, 1))
aux_labelv.data.resize_(batch_size).copy_(
torch.from_numpy(c_label))
fake = self.netG(noisev)
dis_labelv.data.fill_(fake_label)
dis_out, aux_out = self.netD(fake.detach())
dis_errD_fake = dis_criterion(dis_out, dis_labelv)
aux_errD_fake = aux_criterion(aux_out, aux_labelv)
errD_fake = dis_errD_fake + aux_errD_fake
errD_fake.backward()
D_G_z1 = dis_out.data.mean()
errD = errD_real + errD_fake
optimizerD.step()
############################
# (2) Update G network: maximize log(D(G(z)))
###########################
for p in self.netD.parameters():
p.requires_grad = False # to avoid computation
self.netG.zero_grad()
dis_labelv.data.fill_(
real_label) # fake labels are real for generator cost
dis_out, aux_out = self.netD(fake)
dis_errG = dis_criterion(dis_out, dis_labelv)
aux_errG = aux_criterion(aux_out, aux_labelv)
errG = dis_errG + aux_errG
errG.backward()
D_G_z2 = dis_out.data.mean()
optimizerG.step()
print(
'[%d/%d][%d/%d] Loss_D_real: %.4f Loss_D_fake: %.4f Loss_G: %.4f D(x): %.4f D(G(z)): %.4f / %.4f'
% (epoch + 1, self.niter, i + 1, len(
self.data_loader), errD_real.data[0],
errD_fake.data[0], errG.data[0], D_x, D_G_z1, D_G_z2))
if i % 100 == 0:
vutils.save_image(real_cpu,
'%s/real_samples.png' % self.outf,
normalize=True)
fake = self.netG(fixed_noisev)
vutils.save_image(fake.data,
'%s/fake_samples_epoch_%03d.png' %
(self.outf, epoch + 1),
normalize=True)
# do checkpointing
torch.save(self.netG.state_dict(),
'%s/netG_epoch_%03d.pth' % (self.outf, epoch + 1))
torch.save(self.netD.state_dict(),
'%s/netD_epoch_%03d.pth' % (self.outf, epoch + 1))
# compute the average loss
avg_loss_G.update(errG.item(), batch_size)
avg_loss_D.update(errD.item(), batch_size)
avg_loss_A.update(accuracy, batch_size)
avg_loss_M.update(mi.item(), batch_size)
print('[%d/%d][%d/%d] Loss_D: %.4f (%.4f) Loss_G: %.4f (%.4f) D(x): %.4f D(G(z)): %.4f / %.4f Acc: %.4f (%.4f) MI: %.4f (%.4f)'
% (epoch, opt.niter, i, len(dataloader),
errD.item(), avg_loss_D.avg,
errG.item(), avg_loss_G.avg,
D_x, D_G_z1, D_G_z2,
accuracy, avg_loss_A.avg,
mi.item(), avg_loss_M.avg))
if i % 100 == 0:
vutils.save_image(
utils.normalize(real_cpu), '%s/real_samples.png' % opt.outf)
# print('Label for eval = {}'.format(eval_label))
fake = netG(eval_noise, eval_label)
vutils.save_image(
utils.normalize(fake.data),
'%s/fake_samples_epoch_%03d.png' % (opt.outf, epoch)
)
# update eval_label
if opt.visualize_class_label >= 0 and opt.label_rotation:
eval_label_const = (eval_label_const + 1) % num_classes
eval_label.data.fill_(eval_label_const)
# compute metrics
is_mean, is_std, fid = get_metrics(sampler, num_inception_images=opt.num_inception_images, num_splits=10,
prints=True, use_torch=False)
target = torch.full((b_size, ), real_label).to(device)
output = netD(fake_img_tensor)
D_G_z2 = output.mean().item()
errG = criterion(output, target)
errG.backward()
optimizerG.step()
# Save Losses for plotting later
G_losses.append(errG.item())
D_losses.append(errD.item())
if i % 100 == 0:
print(
f'[{epoch+1}/{num_epochs}][{i}/{len(dataloader)}] Loss_D: {errD.item():.4f} Loss_G: {errG.item():.4f}\tD(x): {D_x:.4f}\tD(G(z)): {D_G_z1:.4f} / {D_G_z2:.4f}'
)
vutils.save_image(real_img,
'%s/real_img%03d.jpg' %
("./celeba/load_weight_progress2", epoch + 1),
normalize=True)
with torch.no_grad():
fake = netG(noise).detach()
vutils.save_image(fake.data,
'%s/fake_img_epoch_%03d.jpg' %
("./celeba/load_weight_progress", epoch + 1),
normalize=True)
with torch.no_grad():
fake = netG(fixed_noise).detach()
vutils.save_image(fake.data,
'%s/fake_img_version_epoch_%03d.jpg' %
("./celeba/load_weight_results2", epoch + 1),
normalize=True)
def train(self):
self.train_hist = {}
self.train_hist['D_loss'] = []
self.train_hist['G_loss'] = []
self.train_hist['per_epoch_time'] = []
self.real_labels = torch.ones(self.batch_size, 1).to(device=self.device)
self.fake_labels = torch.zeros(self.batch_size, 1).to(device=self.device)
#Train
print('Training Starts.')
for epoch in range(self.loaded_epoch, self.epochs):
epoch_start_time = time.time()
print('Epoch [%4d/%4d]:' % ((epoch + 1), self.epochs))
for iter, (x, _) in enumerate(self.data_loader): #_ means discard something; y is discarded.
if iter == self.data_loader.dataset.__len__() // self.batch_size:
break
x = x.view(-1, self.data_dim).to(device=self.device) #(-1, 28*28) flatten
z = torch.randn((self.batch_size, self.input_dim)).to(device=self.device)
# 1. Training the Discriminator
# 1-1. Compute D loss
outputs = self.D(x)
D_real_loss = self.BCE_loss(outputs, self.real_labels)
real_score = outputs
outputs = self.D(self.G(z))
D_fake_loss = self.BCE_loss(outputs, self.fake_labels)
D_loss = 0.5 * (D_real_loss + D_fake_loss) # Loss of Discriminator
fake_score = outputs
# 1-2. Backpropagate and Optimize
self.train_hist['D_loss'].append(D_loss.item()) # tensor.item: returns the value of this tensor as a python data.
self.reset_grad()
D_loss.backward() # Calculate gradients.
self.D_optimizer.step() # Performs a single optimization step.
# 2. Training the Generator
# 2.1 Compute G loss
fake_images = self.G(z)
outputs = self.D(fake_images)
G_loss = self.BCE_loss(outputs, self.real_labels)
# 2.2 Backpropagate and Optimize
self.train_hist['G_loss'].append(G_loss.item())
self.reset_grad()
G_loss.backward()
self.G_optimizer.step()
if ((iter + 1) % 200) == 0:
print(" Iter [{}/{}]: D_loss: {:.4f}, G_loss: {:.4f}, D(x): {:.2f}, D(G(z)): {:.2f}"
.format(iter + 1, self.data_loader.dataset.__len__() // self.batch_size, D_loss.item(), G_loss.item(), real_score.mean().item(), fake_score.mean().item()))
self.train_hist['per_epoch_time'].append(time.time() - epoch_start_time)
avg_epoch_time = int(np.mean(self.train_hist['per_epoch_time']))
print(" Avg. 1 epoch time: [%s] / Est. remaining time: [%s]" % (str(datetime.timedelta(seconds = avg_epoch_time)),
str(datetime.timedelta(seconds = (self.epochs - epoch - 1)*avg_epoch_time))))
if (epoch + 1) % self.saving_epoch_interval == 0:
self.save(epoch)
# Save sampled images
fake_images = (fake_images + 1) / 2
fake_images = fake_images.reshape(self.G(z).size(0), 1, 28, 28)
save_image(fake_images, os.path.join(self.sample_dir, 'fake_epoch_{}.png'.format(epoch+1)))
print("Training Finish.")
def save_data(batch, folder, scene_id, batch_idx, save_images, save_txt_depth,
save_binary_depth, save_cam, save_sem_images, save_txt_semantics,
save_binary_semantics):
# Create subfolder using ID of the scene
full_path = os.path.join(folder, scene_id)
print("\nSaving under directory:", full_path)
try:
os.mkdir(full_path)
except FileExistsError:
pass
# Each value in batch dict: List of size n (different views)
# NOTE: n is default 2, focusing on image pairs for now.
img_batch0, img_batch1 = batch["images"] if save_images else (None, None)
depth_batch0, depth_batch1 = batch[
"depths"] if save_txt_depth or save_binary_depth else (None, None)
cam_batch0, cam_batch1 = batch["cameras"] if save_cam else (None, None)
semantic_batch0, semantic_batch1 = batch[
"semantics"] if save_sem_images or save_txt_semantics or save_binary_semantics else (
None, None)
file_prefix = os.path.join(full_path, scene_id)
cam_file_content = "{:<12} = {}';\n"
sample_batch_size = batch["images"][0].shape[0]
start_idx = batch_idx * sample_batch_size
num_views = 2
exts = []
if save_images:
exts.append("png")
if save_txt_depth:
exts.append("depth")
if save_binary_depth:
exts.append("depth.npy")
if save_cam:
exts.append("txt")
if save_sem_images:
exts.append("seg.png")
if save_txt_semantics:
exts.append("semantic")
if save_binary_semantics:
exts.append("semantic.npy")
files_per_view = len(exts)
files_per_sample = num_views * files_per_view
exts = "{" + ",".join(exts) + "}"
for sample_idx in range(sample_batch_size):
curr_file_idx = str(start_idx + sample_idx)
template = file_prefix + "_" + curr_file_idx + "_{pair_id}.{ext}"
# Save RGB images (scene_idx_pairid.png)
if save_images:
rgb_img0, rgb_img1 = img_batch0[sample_idx].cpu(
), img_batch1[sample_idx].cpu()
save_image(rgb_img0, template.format(pair_id=0, ext='png'))
save_image(rgb_img1, template.format(pair_id=1, ext='png'))
# Save depth information
if save_txt_depth or save_binary_depth:
depth0, depth1 = depth_batch0[sample_idx].squeeze(0).cpu().numpy(),\
depth_batch1[sample_idx].squeeze(0).cpu().numpy()
# Save depth information as text file (scene_idx_pairid.depth)
if save_txt_depth:
np.savetxt(template.format(pair_id=0, ext='depth'),
depth0.ravel(),
fmt='%.5f',
delimiter=' ',
newline=' ')
np.savetxt(template.format(pair_id=1, ext='depth'),
depth1.ravel(),
fmt='%.5f',
delimiter=' ',
newline=' ')
# Save depth information as binary file (scene_idx_pairid.depth.npy)
if save_binary_depth:
np.save(template.format(pair_id=0, ext='depth.npy'), depth0)
np.save(template.format(pair_id=1, ext='depth.npy'), depth1)
# Save camera parameters (scene_idx_pairid.txt)
# NOTE: According to SynSin implementation of get_camera_matrices (@camera_transformations.py):
# P: World->Cam, Pinv: Cam->World
if save_cam:
P0, K0, Pinv0, Kinv0 = cam_batch0["P"][sample_idx],\
cam_batch0["K"][sample_idx],\
cam_batch0["Pinv"][sample_idx],\
cam_batch0["Kinv"][sample_idx]
cam_pos, cam_up, cam_dir = split_RT(Pinv0.cpu().numpy())
info = cam_file_content.format("cam_pos", cam_pos)
info += cam_file_content.format("cam_dir", cam_dir)
info += cam_file_content.format("cam_up", cam_up)
with open(template.format(pair_id=0, ext='txt'), 'w+') as f:
f.write(info)
P1, K1, Pinv1, Kinv1 = cam_batch1["P"][sample_idx],\
cam_batch1["K"][sample_idx],\
cam_batch1["Pinv"][sample_idx],\
cam_batch1["Kinv"][sample_idx]
cam_pos, cam_up, cam_dir = split_RT(Pinv1.cpu().numpy())
info = cam_file_content.format("cam_pos", cam_pos)
info += cam_file_content.format("cam_dir", cam_dir)
info += cam_file_content.format("cam_up", cam_up)
with open(template.format(pair_id=1, ext='txt'), 'w+') as f:
f.write(info)
# Save semantic information in the form of int IDs
if save_sem_images or save_txt_semantics or save_binary_semantics:
semantic0, semantic1 = semantic_batch0[sample_idx].squeeze(0).cpu().numpy(),\
semantic_batch1[sample_idx].squeeze(0).cpu().numpy()
if save_sem_images:
sem_img0 = get_semantic_image(semantic0)
sem_img1 = get_semantic_image(semantic1)
sem_img0.save(template.format(pair_id=0, ext='seg.png'))
sem_img1.save(template.format(pair_id=1, ext='seg.png'))
# Save semantic information as text file (scene_idx_pairid.semantic)
if save_txt_semantics:
np.savetxt(template.format(pair_id=0, ext='semantic'),
semantic0.ravel(),
fmt='%d',
delimiter=' ',
newline=' ')
np.savetxt(template.format(pair_id=1, ext='semantic'),
semantic1.ravel(),
fmt='%d',
delimiter=' ',
newline=' ')
# Save semantic information as binary file (scene_idx_pairid.semantic.npy)
if save_binary_semantics:
np.save(template.format(pair_id=0, ext='semantic.npy'),
semantic0)
np.save(template.format(pair_id=1, ext='semantic.npy'),
semantic1)
print("Files created: {}_{}_{{0,1}}.{}".format(scene_id, curr_file_idx,
exts))
print("Saving completed. Number of files created under {}: {}\n".format(
full_path, files_per_sample * sample_batch_size))
def main(opts):
# Create the data loader
dataset = DatasetFromFolder(root, 'train')
loader = DataLoader(dataset, batch_size=opts.batch_size, shuffle=True)
# Create the model
start_epoch = 0
G = StyleGenerator()
D = StyleDiscriminator()
# Load the pre-trained weight
if os.path.exists(opts.checkpoint):
print("Load checkpoint")
state = torch.load(opts.checkpoint)
G.load_state_dict(state['G'])
D.load_state_dict(state['D'])
start_epoch = state['start_epoch']
else:
print(
"Pre-trained weight cannot load successfully, train from scratch!")
# Multi-GPU support
if torch.cuda.device_count() > 1:
print("Multiple GPU:" + str(torch.cuda.device_count()) + "\t GPUs")
G = nn.DataParallel(G)
D = nn.DataParallel(D)
G.to(opts.device)
D.to(opts.device)
# Create the criterion, optimizer and scheduler
optim_D = optim.Adam(D.parameters(), lr=0.00001, betas=(0.5, 0.999))
optim_G = optim.Adam(G.parameters(), lr=0.00001, betas=(0.5, 0.999))
scheduler_D = optim.lr_scheduler.ExponentialLR(optim_D, gamma=0.99)
scheduler_G = optim.lr_scheduler.ExponentialLR(optim_G, gamma=0.99)
# Train
fix_z = torch.randn([opts.batch_size, 512]).to(opts.device)
softplus = nn.Softplus()
Loss_D_list = [0.0]
Loss_G_list = [0.0]
for ep in range(start_epoch, opts.epoch):
bar = tqdm(loader)
loss_D_list = []
loss_G_list = []
for i, (real_img, ) in enumerate(bar):
# =======================================================================================================
# (1) Update D network: maximize log(D(x)) + log(1 - D(G(z)))
# =======================================================================================================
# Compute adversarial loss toward discriminator
D.zero_grad()
real_img = real_img.to(opts.device)
real_logit = D(real_img)
fake_img = G(torch.randn([real_img.size(0), 512]).to(opts.device))
fake_logit = D(fake_img.detach())
d_loss = softplus(fake_logit).mean()
d_loss = d_loss + softplus(-real_logit).mean()
if opts.r1_gamma != 0.0:
r1_penalty = R1Penalty(real_img.detach(), D)
d_loss = d_loss + r1_penalty * (opts.r1_gamma * 0.5)
if opts.r2_gamma != 0.0:
r2_penalty = R2Penalty(fake_img.detach(), D)
d_loss = d_loss + r2_penalty * (opts.r2_gamma * 0.5)
loss_D_list.append(d_loss.item())
# Update discriminator
d_loss.backward()
optim_D.step()
# =======================================================================================================
# (2) Update G network: maximize log(D(G(z)))
# =======================================================================================================
if i % CRITIC_ITER == 0:
G.zero_grad()
fake_logit = D(fake_img)
g_loss = softplus(-fake_logit).mean()
loss_G_list.append(g_loss.item())
# Update generator
g_loss.backward()
optim_G.step()
# Output training stats
bar.set_description("Epoch {} [{}, {}] [G]: {} [D]: {}".format(
ep, i + 1, len(loader), loss_G_list[-1], loss_D_list[-1]))
# Save the result
Loss_G_list.append(np.mean(loss_G_list))
Loss_D_list.append(np.mean(loss_D_list))
# Check how the generator is doing by saving G's output on fixed_noise
with torch.no_grad():
fake_img = G(fix_z).detach().cpu()
save_image(fake_img,
os.path.join(opts.det, 'images',
str(ep) + '.png'),
nrow=4,
normalize=True)
# Save model
state = {
'G': G.state_dict(),
'D': D.state_dict(),
'Loss_G': Loss_G_list,
'Loss_D': Loss_D_list,
'start_epoch': ep,
}
torch.save(state, os.path.join(opts.det, 'models', 'latest.pth'))
scheduler_D.step()
scheduler_G.step()
# Plot the total loss curve
Loss_D_list = Loss_D_list[1:]
Loss_G_list = Loss_G_list[1:]
plotLossCurve(opts, Loss_D_list, Loss_G_list)
recon_loss.data.item()))
if (epoch + 1) % 1 == 0:
batch_size = 104
test_iter = iter(test_loader)
test_data = next(test_iter)
z_real = encoder(Variable(test_data[0]).cuda())
reconst = decoder(torch.randn_like(z_real)).cpu().view(
batch_size, 1, 28, 28)
if not os.path.isdir('./data/reconst_images'):
os.makedirs('data/reconst_images')
save_image(
test_data[0].view(batch_size, 1, 28, 28),
'./data/reconst_images/wae_gan_clamp_input_%d.png' %
(epoch + 1))
save_image(
reconst.data,
'./data/reconst_images/wae_gan_clamp_images_%d.png' %
(epoch + 1))
filename = "wae_gan_mnist_clamp"
torch.save(encoder.state_dict(), filename + "_enc")
torch.save(decoder.state_dict(), filename + "_dec")
def load_model(filename):
encoder = Encoder()
decoder = Decoder()
encoder.load_state_dict(torch.load(filename + "_enc"))
def save_pic(x_rec):
resultsample = x_rec * 0.5 + 0.5
resultsample = resultsample.cpu()
save_image(resultsample, 'sample_%i_lr.png' % i, nrow=16)
input = Variable(input)
if args.use_cuda:
input = input.cuda()
optimizer.zero_grad()
nll, kld = model.loss(input, criterion, eval=False)
loss = nll + kld
loss.backward()
optimizer.step()
if iteration % 10 == 0:
print('epoch {}, iteration {}, loss {}'.format(
epoch, iteration,
loss.cpu().data.numpy()[0]))
print(F.softmax(model.p_c_logits, dim=0).cpu().data.numpy())
sampling = model.sample(args.use_cuda)
vutils.save_image(
sampling,
'sampling/vae_{}.png'.format(epoch * len(dataloader) +
iteration))
samplings = [f for f in listdir('sampling')]
samplings = [
imageio.imread('sampling/' + path) for path in samplings
for _ in range(2)
]
imageio.mimsave('sampling/movie.gif', samplings)
avg_rec_loss = 0
avg_kld = 0
for i, (img1, text1, len_text1, img2, text2, len_text2) in tqdm(enumerate(train_loader)):
img1, text_feat1, text_feat1_mismatch, img2, text_feat2, text_feat2_mismatch = \
preprecessing(img1, text1, len_text1, img2, text2, len_text2, text_encoder)
##########################################################
#### one for discrimator ####
img1_D = img1
img2_D = img2
##########################################################
if epoch % 20 == 0 and i == 0:
save_image(img1.data, './result/origin_epoch_%d.jpg' % (epoch + 1))
save_image(img2.data, './result/target_epoch_%d.jpg' % (epoch + 1))
ONES = Variable(torch.ones(img1.size(0))) # size(0) is the number of the minibatch samples
ZEROS = Variable(torch.zeros(img1.size(0)))
if config.is_cuda:
ONES, ZEROS = ONES.cuda(), ZEROS.cuda()
##############################################
### UPDATE DISCRIMINATOR ###
##############################################
# ------------- img1, text1 pairs ------------
D.zero_grad()
# real image with matching text
real_logit = D(img1_D, text_feat1)
size = 256
input_A = torch.ones([1, 3, size, size], dtype=torch.float).to(device)
input_B = torch.ones([1, 3, size, size], dtype=torch.float).to(device)
transforms_ = [
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
]
data_root = "/home/kuan/workspace/muke/pytorch/09/datasets/apple2orange"
dataloader = DataLoader(ImageDataset(data_root, transforms_),
batch_size=1,
shuffle=False,
num_workers=8)
if not os.path.exists("outputs/A"):
os.makedirs("outputs/A")
if not os.path.exists("outputs/B"):
os.makedirs("outputs/B")
for i, batch in enumerate(dataloader):
real_A = torch.tensor(input_A.copy_(batch['A']),
dtype=torch.float).to(device)
real_B = torch.tensor(input_B.copy_(batch['B']),
dtype=torch.float).to(device)
fake_B = 0.5 * (netG_A2B(real_A).data + 1.0)
fake_A = 0.5 * (netG_B2A(real_B).data + 1.0)
save_image(fake_A, "outputs/A/{}.png".format(i))
save_image(fake_B, "outputs/B/{}.png".format(i))
print(i)
], 0).cuda()
hint = torch.cat((real_vim * mask, mask), 1)
if flag: # fix samples
writer.add_image(
'target imgs',
vutils.make_grid(real_cim.mul(0.5).add(0.5), nrow=16))
writer.add_image(
'sketch imgs',
vutils.make_grid(real_sim.mul(0.5).add(0.5), nrow=16))
writer.add_image(
'hint',
vutils.make_grid((real_vim * mask).mul(0.5).add(0.5),
nrow=16))
vutils.save_image(
real_cim.mul(0.5).add(0.5),
'%s/color_samples' % opt.outf + '.png')
vutils.save_image(
real_sim.mul(0.5).add(0.5),
'%s/blur_samples' % opt.outf + '.png')
fixed_sketch.resize_as_(real_sim).copy_(real_sim)
fixed_hint.resize_as_(hint).copy_(hint)
flag -= 1
fake = netG(Variable(real_sim), Variable(hint))
if gen_iterations < opt.baseGeni:
contentLoss = criterion_L2(
netF((fake.mul(0.5) - Variable(saber)) / Variable(diver)),
netF(Variable((real_cim.mul(0.5) - saber) / diver)))
#======================================================================
for step in tqdm(range(args.n_samplings), desc="Samplings"):
model_G.eval()
model_D.eval()
# seed 値の変更
np.random.seed(seed)
torch.manual_seed(seed)
# 入力ノイズ z の再生成
input_noize_z = torch.randn(size=(args.batch_size,
args.n_input_noize_z, 1,
1)).to(device)
#====================================================
# 生成器 G の 推論処理
#====================================================
with torch.no_grad():
# G(z) : 生成器から出力される偽物画像
G_z = model_G(input_noize_z)
#---------------------
# 出力画像の生成＆保存
#---------------------
save_image(tensor=G_z,
filename=os.path.join(args.results_dir, args.exper_name) +
"/fake_image_seed{}_batchAll.png".format(seed))
seed += 1
n_print -= 1
print("Finished Test Loop.")
#save_image(x, './mlp_img/k_values.eps', nrow=10)
# Decode k values in order
node_cluster = np.zeros(128).astype(np.float32)
for idx in range(10):
count = 0
for i in range(K):
if c_label[i] == idx and count < 5:
node_cluster = np.c_[node_cluster, c_cluster[i, :]]
count += 1
node_cluster = np.transpose(node_cluster)
node_cluster = np.delete(node_cluster, (0), axis=0)
x = Variable(torch.from_numpy(node_cluster)).cuda()
x = to_img(model[1](x).cpu().data)
save_image(x, './mlp_img/k_values_order.eps', nrow=10)
#-------------------------------------------------
# Find images that mapped to cluster point
#-------------------------------------------------
# Find images mapped to each cluster
for j in range(K):
count = 0
cluster_i = np.zeros(128).astype(np.float32)
for i in range(dataset_size):
if cl[i] == j and count < 8:
cluster_i = np.c_[cluster_i, Feat[i]]
count += 1
cluster_i = np.transpose(cluster_i)
cluster_i = np.delete(cluster_i, (0), axis=0)
x_cluster_i = Variable(torch.from_numpy(cluster_i)).cuda()
for batch_idx, data in enumerate(dataloader):
img, _ = data
num_img = img.size(0)
img = img.view(num_img, -1)
img = Variable(img)
if torch.cuda.is_available():
img = img.cuda()
# forward pass
recon_out, mu, logvar = model(img)
loss = loss_function(recon_out, img, mu, logvar)
# backward
optimizer.zero_grad()
loss.backward()
optimizer.step()
train_loss += loss.data
if batch_idx % 100 == 0:
print('Epoch:{} [{}/{} ({:.0f}%)]\tloss:{:.6f}'.format(
epoch + 1, batch_idx * len(img), len(dataloader.dataset),
100.0 * batch_idx / len(dataloader), loss.data[0] / len(img)))
print('Epoch {} Finished, Average loss:{:.4f}'.format(
epoch + 1, train_loss / len(dataloader.dataset)))
if epoch % 10 == 0:
pic = to_img(recon_out.cpu().data)
save_image(pic, './vae_img/images_{}.png'.format(epoch))
torch.save(model.state_dict(), './variational_autoencoder.pth')
g = g.to(device)
z = np.random.RandomState(0).randn(n_sample, 512).astype("float32")
with torch.no_grad():
img_pt, _ = g(
[torch.from_numpy(z).to(device)],
truncation=0.5,
truncation_latent=latent_avg.to(device),
randomize_noise=False,
)
Gs_kwargs = dnnlib.EasyDict()
Gs_kwargs.randomize_noise = False
img_tf = g_ema.run(z, None, **Gs_kwargs)
img_tf = torch.from_numpy(img_tf).to(device)
img_diff = ((img_pt + 1) / 2).clamp(0.0, 1.0) - (
(img_tf.to(device) + 1) / 2).clamp(0.0, 1.0)
img_concat = torch.cat((img_tf, img_pt, img_diff), dim=0)
print(img_diff.abs().max())
utils.save_image(img_concat,
name + ".png",
nrow=n_sample,
normalize=True,
range=(-1, 1))
def sample(cfg, logger):
torch.cuda.set_device(0)
model = SoftIntroVAEModelTL(
startf=cfg.MODEL.START_CHANNEL_COUNT,
layer_count=cfg.MODEL.LAYER_COUNT,
maxf=cfg.MODEL.MAX_CHANNEL_COUNT,
latent_size=cfg.MODEL.LATENT_SPACE_SIZE,
dlatent_avg_beta=cfg.MODEL.DLATENT_AVG_BETA,
style_mixing_prob=cfg.MODEL.STYLE_MIXING_PROB,
mapping_layers=cfg.MODEL.MAPPING_LAYERS,
channels=cfg.MODEL.CHANNELS,
generator=cfg.MODEL.GENERATOR,
encoder=cfg.MODEL.ENCODER,
beta_kl=cfg.MODEL.BETA_KL,
beta_rec=cfg.MODEL.BETA_REC,
beta_neg=cfg.MODEL.BETA_NEG[cfg.MODEL.LAYER_COUNT - 1],
scale=cfg.MODEL.SCALE)
model.cuda(0)
model.eval()
model.requires_grad_(False)
decoder = model.decoder
encoder = model.encoder
mapping_tl = model.mapping_tl
mapping_fl = model.mapping_fl
dlatent_avg = model.dlatent_avg
logger.info("Trainable parameters decoder:")
print(count_parameters(decoder))
logger.info("Trainable parameters encoder:")
print(count_parameters(encoder))
arguments = dict()
arguments["iteration"] = 0
model_dict = {
'discriminator_s': encoder,
'generator_s': decoder,
'mapping_tl_s': mapping_tl,
'mapping_fl_s': mapping_fl,
'dlatent_avg': dlatent_avg
}
checkpointer = Checkpointer(cfg, model_dict, {}, logger=logger, save=False)
extra_checkpoint_data = checkpointer.load()
model.eval()
layer_count = cfg.MODEL.LAYER_COUNT
def encode(x):
z, mu, _ = model.encode(x, layer_count - 1, 1)
styles = model.mapping_fl(mu)
return styles
def decode(x):
return model.decoder(x, layer_count - 1, 1, noise=True)
path = cfg.DATASET.SAMPLES_PATH
im_size = 2**(cfg.MODEL.LAYER_COUNT + 1)
paths = list(os.listdir(path))
paths = sorted(paths)
random.seed(5)
random.shuffle(paths)
def move_to(list, item, new_index):
list.remove(item)
list.insert(new_index, item)
def make(paths):
src = []
for filename in paths:
img = np.asarray(Image.open(path + '/' + filename))
if img.shape[2] == 4:
img = img[:, :, :3]
im = img.transpose((2, 0, 1))
x = torch.tensor(np.asarray(im, dtype=np.float32),
requires_grad=True).cuda() / 127.5 - 1.
if x.shape[0] == 4:
x = x[:3]
factor = x.shape[2] // im_size
if factor != 1:
x = torch.nn.functional.avg_pool2d(x[None, ...], factor,
factor)[0]
assert x.shape[2] == im_size
src.append(x)
with torch.no_grad():
reconstructions = []
for s in src:
latents = encode(s[None, ...])
reconstructions.append(decode(latents).cpu().detach().numpy())
return src, reconstructions
def chunker_list(seq, size):
return list((seq[i::size] for i in range(size)))
final = chunker_list(paths, 4)
path0, path1, path2, path3 = final
path0.reverse()
path1.reverse()
path2.reverse()
path3.reverse()
src0, rec0 = make(path0)
src1, rec1 = make(path1)
src2, rec2 = make(path2)
src3, rec3 = make(path3)
initial_resolution = im_size
lods_down = 1
padding_step = 4
width = 0
height = 0
current_padding = 0
final_resolution = initial_resolution
for _ in range(lods_down):
final_resolution /= 2
for i in range(lods_down + 1):
width += current_padding * 2**(lods_down - i)
height += current_padding * 2**(lods_down - i)
current_padding += padding_step
width += 2**(lods_down + 1) * final_resolution
height += (lods_down + 1) * initial_resolution
width = int(width)
height = int(height)
def make_part(current_padding, src, rec):
canvas = np.ones([3, height + 20, width + 10])
padd = 0
initial_padding = current_padding
height_padding = 0
for i in range(lods_down + 1):
for x in range(2**i):
for y in range(2**i):
try:
ims = src.pop()
imr = rec.pop()[0]
ims = ims.cpu().detach().numpy()
imr = imr
res = int(initial_resolution / 2**i)
ims = resize(ims, (3, initial_resolution / 2**i,
initial_resolution / 2**i))
imr = resize(imr, (3, initial_resolution / 2**i,
initial_resolution / 2**i))
place(
canvas, ims,
current_padding + x * (2 * res + current_padding),
i * initial_resolution + height_padding + y *
(res + current_padding))
place(
canvas, imr, current_padding + res + x *
(2 * res + current_padding),
i * initial_resolution + height_padding + y *
(res + current_padding))
except IndexError:
return canvas
height_padding += initial_padding * 2
current_padding -= padding_step
padd += padding_step
return canvas
canvas = [
make_part(current_padding, src0, rec0),
make_part(current_padding, src1, rec1),
make_part(current_padding, src2, rec2),
make_part(current_padding, src3, rec3)
]
canvas = np.concatenate(canvas, axis=2)
path = './make_figures/output'
os.makedirs(path, exist_ok=True)
os.makedirs(os.path.join(path, cfg.NAME), exist_ok=True)
print('Saving image')
save_path = './make_figures/output/%s/reconstructions_multiresolution.png' % cfg.NAME
os.makedirs(os.path.dirname(save_path), exist_ok=True)
save_image(torch.Tensor(canvas), save_path)
npImg1 = cv2.imread("HRS_99H8_4801_s144.bmp")
img1 = torch.from_numpy(np.rollaxis(npImg1, 2)).float().unsqueeze(0)
img2 = torch.randn(img1.size(), requires_grad = True)
print(img1.size())
print(img2.size())
# device = torch.device('cuda')
# if torch.cuda.is_available():
# img1 = img1.to(device)
# img2 = img2.to(device)
criterion = ML1Loss(img1.size()[2], img1.size()[3])
optimizer = optim.Adam([img2], lr=0.01)
count = 0
loss = 1000
while loss > 100:
optimizer.zero_grad()
loss = criterion(img2, img1, ["HRS_99H8_4801_s144.bmp"])
count += 1
# print("{0}: {1}".format(count,loss))
loss.backward()
optimizer.step()
utils.save_image(img2,"img2.png",normalize=True, range=(0,255))
# utils.save_image(img1,"img1.png",normalize=True, range=(0,255))
def train(self):
self.create_model()
losses = {'perceptual': [], 'content': []}
optimizer = torch.optim.Adam(itertools.chain(
self.attention1.parameters(), self.encoder.parameters(),
self.decoder.parameters()),
lr=self.config.learning_rate)
criterion = torch.nn.L1Loss()
criterion_p = torch.nn.MSELoss(reduction='mean')
styles = iter(self.style_loader)
self.encoder.train()
self.decoder.train()
self.reporter.writeInfo("Start to train the model")
for e in range(1, self.config.epoch_size + 1):
for i, (content, target) in enumerate(self.train_loader):
try:
style, target = next(styles)
except:
styles = iter(self.style_loader)
style, target = next(styles)
content = content.cuda()
style = style.cuda()
fea_c = self.encoder(content)
fea_s = self.encoder(style)
out_feature, attention_map = self.attention1(fea_c, fea_s)
out_content = self.decoder(out_feature)
c1, c2, c3, _ = self.D(content)
h1, h2, h3, _ = self.D(out_content)
s1, s2, s3, _ = self.D(style)
loss_content = torch.norm(c3 - h3, p=2)
loss_perceptual = 0
for t in range(3):
loss_perceptual += criterion(
gram_matrix(eval('s' + str(t + 1))),
gram_matrix(eval('h' + str(t + 1))))
loss = loss_content * self.config.content_weight + loss_perceptual * self.config.style_weight
optimizer.zero_grad()
loss.backward()
optimizer.step()
if i % self.config.log_interval == 0:
now = datetime.datetime.now()
otherStyleTime = now.strftime("%Y-%m-%d %H:%M:%S")
print(otherStyleTime)
print('epoch: ', e, ' iter: ', i)
print(
'attention scartters: ',
torch.std(attention_map.argmax(-1).float(),
1).mean().cpu())
print(attention_map.shape)
self.attention1.hard = True
self.attention1.eval()
self.decoder.eval()
tosave, perc, cont = self.eval()
save_image(
denorm(tosave),
self.image_dir + '/epoch_{}-iter_{}.png'.format(e, i))
print("image saved to " + self.image_dir +
'/epoch_{}-iter_{}.png'.format(e, i))
print('content loss:', cont)
print('perceptual loss:', perc)
self.reporter.writeTrainLog(
e, i, f'''
attention scartters: {torch.std(attention_map.argmax(-1).float(), 1).mean().cpu()}\n
content loss: {cont}\n
perceptual loss: {perc}
''')
losses['perceptual'].append(perc)
losses['content'].append(cont)
self.plot_loss_curve(losses)
self.attention1.hard = False
self.attention1.train()
self.decoder.train()
torch.save(
{
'layer1': self.attention1.state_dict(),
'encoder': self.encoder.state_dict(),
'decoder': self.decoder.state_dict()
}, f'{self.model_state_dir}/epoch_{e}-iter_{i}.pth')
print('model saved.')
