def main(args):
domain = args.domain
batch_size = args.batch_size
nets = networks.define_nets(args.model,
args.domain,
load_encoder=True,
device='cuda')
zs = nets.sample_zs(args.num_samples, args.seed)
for batch_start in pbar(range(0, args.num_samples, batch_size)):
s = slice(batch_start, min(batch_start + batch_size, args.num_samples))
zs_batch = zs[s]
with torch.no_grad():
ims_initial = nets.zs2image(zs_batch)
ims = nets.invert(ims_initial, mask=None)
for i, im in enumerate(ims):
filename = os.path.join(
args.outdir,
'seed%03d_sample%05d' % (args.seed, i + batch_start))
pil_image = renormalize.as_image(im)
if args.im_size:
pil_image = pil_image.resize((args.im_size, args.im_size),
Image.ANTIALIAS)
pil_image.save(filename + '.png')
def save_zds_images(dirname, model, zds,
name_template="image_{}.png", batch_size=10, indices=None):
if indices is not None:
class Sampler(torch.utils.data.Sampler):
def __init__(self):
pass
def __iter__(self):
yield from indices
def __len__(self):
return len(indices)
sampler = Sampler()
else:
sampler = None
os.makedirs(dirname, exist_ok=True)
with torch.no_grad():
# Now generate images
z_loader = torch.utils.data.DataLoader(IndexDataset(zds),
batch_size=batch_size, num_workers=2, sampler=sampler,
pin_memory=True)
saver = imgsave.SaveImagePool()
for indices, [z] in pbar(z_loader, desc='Saving images'):
z = z.cuda()
im = model(z).cpu()
for i, index in enumerate(indices.tolist()):
filename = os.path.join(dirname, name_template.format(index))
saver.add(renormalize.as_image(im[i]), filename)
saver.join()
def train(self, dataset):
z = tf.constant(random.normal((FLAGS.n_samples, 1, 1, self.z_dim)))
g_train_loss = metrics.Mean()
d_train_loss = metrics.Mean()
for epoch in range(self.epochs):
bar = pbar(self.total_images, self.batch_size, epoch, self.epochs)
for batch in dataset:
for _ in range(self.n_critic):
self.train_d(batch)
d_loss = self.train_d(batch)
d_train_loss(d_loss)
g_loss = self.train_g()
g_train_loss(g_loss)
self.train_g()
bar.postfix['g_loss'] = f'{g_train_loss.result():6.3f}'
bar.postfix['d_loss'] = f'{d_train_loss.result():6.3f}'
bar.update(self.batch_size)
g_train_loss.reset_states()
d_train_loss.reset_states()
bar.close()
del bar
samples = self.generate_samples(z)
image_grid = img_merge(samples, n_rows=8).squeeze()
save_image_grid(image_grid, epoch + 1)
def load_or_compute_features(path, args):
saved_features = os.path.join(path, 'metrics/prdc/features.npz')
if os.path.isfile(saved_features):
print("found saved features: %s" % saved_features)
return np.load(saved_features)['features']
transform = transforms.Compose([
transforms.Resize(args.load_size, Image.ANTIALIAS),
transforms.CenterCrop(args.load_size),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
])
# uses imagenet normalization, with pretrained vgg features
dset = parallelfolder.ParallelImageFolders([path],
transform=transform,
lazy_init=False)
loader = DataLoader(dset,
batch_size=args.batch_size,
shuffle=False,
pin_memory=False,
num_workers=args.workers)
model = features.vgg16().eval().cuda()
feature_list = []
for data in pbar(loader):
data = data[0].cuda()
with torch.no_grad():
feature_list.append(model(data).cpu().numpy())
feature_list = np.concatenate(feature_list)
os.makedirs(os.path.join(path, 'metrics/prdc'), exist_ok=True)
np.savez(saved_features, features=feature_list)
return feature_list
def main(args):
# make all the directories
os.makedirs(osp.join(args.outdir, 'composite_original'), exist_ok=True)
os.makedirs(osp.join(args.outdir, 'composite_mask'), exist_ok=True)
os.makedirs(osp.join(args.outdir, 'inverted_RGBM'), exist_ok=True)
os.makedirs(osp.join(args.outdir, 'poisson'), exist_ok=True)
model_type = args.model
domain = args.domain
nets_RGBM = networks.define_nets(model_type, domain)
compositer = compositions.get_compositer(domain)(nets_RGBM)
if args.input_source == 'images':
assert (args.domain in args.data_path) # sanity check!
outdim = nets_RGBM.setting['outdim']
transform = transforms.Compose([
transforms.Resize(outdim),
transforms.CenterCrop(outdim),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
dset = parallelfolder.ParallelImageFolders([args.data_path],
transform=transform,
lazy_init=False)
print("Using dataset of length: %d" % len(dset))
def resize2image(tensor, method=Image.ANTIALIAS):
return (renormalize.as_image(tensor[0]).resize(
(args.im_size, args.im_size), method))
for i in pbar(range(args.num_samples)):
with torch.no_grad():
if args.input_source == 'samples':
rng = np.random.RandomState(i)
indices = rng.choice(compositer.total_samples,
len(compositer.ordered_labels))
composite_data = compositer(indices)
elif args.input_source == 'images':
rng = np.random.RandomState(i)
indices = rng.choice(len(dset), len(compositer.ordered_labels))
images = [dset[i][0] for i in indices]
composite_data = compositer(indices=None, imgs=images)
resize2image(composite_data.composite_image).save(
os.path.join(args.outdir, 'composite_original',
'sample%06d_composite_original.png' % i))
resize2image(composite_data.composite_mask,
method=Image.NEAREST).save(
os.path.join(args.outdir, 'composite_mask',
'sample%06d_composite_mask.png' % i))
resize2image(composite_data.inverted_RGBM).save(
os.path.join(args.outdir, 'inverted_RGBM',
'sample%06d_inverted_RGBM.png' % i))
poisson = compositions.poisson_blend_layers(
composite_data.parts_image, composite_data.parts_mask)
poisson = poisson.resize((args.im_size, args.im_size),
Image.ANTIALIAS)
poisson.save(
osp.join(args.outdir, 'poisson', 'sample%06d_poisson.png' % i))
def train(opt):
print("Random Seed: ", opt.seed)
random.seed(opt.seed)
torch.manual_seed(opt.seed)
cudnn.benchmark = True
device = 'cuda'
batch_size = int(opt.batch_size)
# tensorboard
os.makedirs(os.path.join(opt.outf, 'runs'), exist_ok=True)
writer = SummaryWriter(log_dir=os.path.join(opt.outf, 'runs'))
# classifier follows architecture and initialization from attribute
# classifiers in stylegan:
# https://github.com/NVlabs/stylegan/blob/master/metrics/linear_separability.py#L136
# https://github.com/NVlabs/stylegan/blob/master/training/networks_stylegan.py#L564
net = attribute_classifier.D(3,
resolution=256,
fixed_size=True,
use_mbstd=False)
# use random normal bc wscale will rescale weights, see tf source
# https://github.com/NVlabs/stylegan/blob/master/training/networks_stylegan.py#L148
netinit.init_weights(net, init_type='normal', gain=1.)
net = net.to(device)
# losses + optimizers
bce_loss = nn.BCEWithLogitsLoss()
optimizer = optim.Adam(net.parameters(),
lr=opt.lr,
betas=(opt.beta1, 0.999))
# datasets -- added random horizonal flipping for training
train_transform = data.get_transform('celebahq', 'imtrain')
train_dset = data.get_dataset('celebahq',
'train',
opt.attribute,
load_w=False,
transform=train_transform)
print("Training transform:")
print(train_transform)
val_transform = data.get_transform('celebahq', 'imval')
val_dset = data.get_dataset('celebahq',
'val',
opt.attribute,
load_w=False,
transform=val_transform)
print("Validation transform:")
print(val_transform)
train_loader = DataLoader(train_dset,
batch_size=opt.batch_size,
shuffle=True,
pin_memory=False,
num_workers=opt.workers)
val_loader = DataLoader(val_dset,
batch_size=opt.batch_size,
shuffle=False,
pin_memory=False,
num_workers=opt.workers)
start_ep = 0
best_val_acc = 0.0
best_val_epoch = 0
for epoch in pbar(range(start_ep, opt.niter + 1)):
# average meter for train/val loss and train/val acc
metrics = dict(train_loss=util.AverageMeter(),
val_loss=util.AverageMeter(),
train_acc=util.AverageMeter(),
val_acc=util.AverageMeter())
# train loop
for step, (im, label) in enumerate(pbar(train_loader)):
im = im.cuda()
label = label.cuda().float()
net.zero_grad()
logit, softmaxed = attribute_utils.get_softmaxed(net, im)
# enforces that negative logit --> our label = 1
loss = bce_loss(logit, 1 - label)
predicted = (softmaxed > 0.5).long()
correct = (predicted == label).float().mean().item()
metrics['train_loss'].update(loss, n=len(label))
metrics['train_acc'].update(correct, n=len(label))
loss.backward()
optimizer.step()
if step % 200 == 0:
pbar.print("%s: %0.2f" %
('train loss', metrics['train_loss'].avg))
pbar.print("%s: %0.2f" %
('train acc', metrics['train_acc'].avg))
# val loop
net = net.eval()
with torch.no_grad():
for step, (im, label) in enumerate(pbar(val_loader)):
im = im.cuda()
label = label.cuda().float()
logit, softmaxed = attribute_utils.get_softmaxed(net, im)
predicted = (softmaxed > 0.5).long()
correct = (predicted == label).float().mean().item()
loss = bce_loss(logit, 1 - label)
metrics['val_loss'].update(loss, n=len(label))
metrics['val_acc'].update(correct, n=len(label))
net = net.train()
# send losses to tensorboard
for k, v in metrics.items():
pbar.print("Metrics at end of epoch")
pbar.print("%s: %0.4f" % (k, v.avg))
writer.add_scalar(k.replace('_', '/'), v.avg, epoch)
# do checkpoint as latest
util.make_checkpoint(net, optimizer, epoch, metrics['val_acc'].avg,
opt.outf, 'latest')
if metrics['val_acc'].avg > best_val_acc:
pbar.print("Updating best checkpoint at epoch %d" % epoch)
pbar.print("Old Best Epoch %d Best Val %0.2f" %
(best_val_epoch, best_val_acc))
# do checkpoint as best
util.make_checkpoint(net, optimizer, epoch, metrics['val_acc'].avg,
opt.outf, 'best')
best_val_acc = metrics['val_acc'].avg
best_val_epoch = epoch
pbar.print("New Best Epoch %d Best Val %0.2f" %
(best_val_epoch, best_val_acc))
with open("%s/best_val.txt" % opt.outf, "w") as f:
f.write("Best Epoch %d Best Val %0.2f\n" %
(best_val_epoch, best_val_acc))
if epoch >= best_val_epoch + 5:
pbar.print("Exiting training")
pbar.print("Best Val epoch %d" % best_val_epoch)
pbar.print("Curr epoch %d" % epoch)
break
def compute_distances(args):
transform = transforms.Compose([
transforms.Resize(args.load_size, Image.ANTIALIAS),
transforms.CenterCrop(args.load_size),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
rec_path = args.path[0].rstrip('/')
if 'car' in rec_path:
assert(args.crop_aspect_car) # sanity check
gt_path = os.path.join(os.path.dirname(rec_path),'composite_original')
mask_path = os.path.join(os.path.dirname(rec_path),'composite_mask')
dset_GT = parallelfolder.ParallelImageFolders([gt_path],
transform=transform,
lazy_init=False,
return_path=True)
dset_mask = parallelfolder.ParallelImageFolders([mask_path],
transform=transform,
lazy_init=False,
return_path=True)
dset_rec = parallelfolder.ParallelImageFolders([rec_path],
transform=transform,
lazy_init=False,
return_path=True)
gt_loader = DataLoader(dset_GT, batch_size=args.batch_size, shuffle=False,
pin_memory=False, num_workers=2)
mask_loader = DataLoader(dset_mask, batch_size=args.batch_size, shuffle=False,
pin_memory=False, num_workers=2)
rec_loader = DataLoader(dset_rec, batch_size=args.batch_size, shuffle=False,
pin_memory=False, num_workers=2)
l1_distance = losses.Masked_L1_Loss()
lpips_alex = losses.Masked_LPIPS_Loss(net='alex')
lpips_vgg = losses.Masked_LPIPS_Loss(net='vgg')
loss_metrics = dict(l1=[], lpips_alex=[], lpips_vgg=[])
for (gt, mask, rec) in zip(pbar(gt_loader), mask_loader, rec_loader):
# check the paths
for p1, p2, p3 in zip(gt[1][0], mask[1][0], rec[1][0]):
sample = p1.split('/')[-1].split('_')[0]
assert(p2.split('/')[-1].split('_')[0] == sample)
assert(p3.split('/')[-1].split('_')[0] == sample)
with torch.no_grad():
gt_tensor = gt[0][0].cuda()
rec_tensor = rec[0][0].cuda()
mask_tensor = mask[0][0] # 0 to 1
mask_tensor[mask_tensor < 0.5] = 0
assert(torch.max(mask_tensor) == 1)
assert(torch.min(mask_tensor) == 0)
if args.crop_aspect_car:
# don't compute similarity on black padding
aspect_border = int(0.125 * args.load_size)
mask_tensor[:, :, :aspect_border, :] = 0
mask_tensor[:, :, -aspect_border:, :] = 0
mask_tensor = mask_tensor[:, :1, :, :].cuda()
batch_l1 = l1_distance(rec_tensor, gt_tensor,
mask_tensor).cpu().numpy().squeeze()
batch_lpips_alex = lpips_alex(rec_tensor, gt_tensor,
mask_tensor).cpu().numpy().squeeze()
batch_lpips_vgg = lpips_vgg(rec_tensor, gt_tensor,
mask_tensor).cpu().numpy().squeeze()
loss_metrics['l1'].append(batch_l1)
loss_metrics['lpips_alex'].append(batch_lpips_alex)
loss_metrics['lpips_vgg'].append(batch_lpips_vgg)
loss_metrics['l1'] = np.concatenate(loss_metrics['l1'])
loss_metrics['lpips_alex'] = np.concatenate(loss_metrics['lpips_alex'])
loss_metrics['lpips_vgg'] = np.concatenate(loss_metrics['lpips_vgg'])
loss_metrics['l1_avg'] = np.mean(loss_metrics['l1'])
loss_metrics['lpips_alex_avg'] = np.mean(loss_metrics['lpips_alex'])
loss_metrics['lpips_vgg_avg'] = np.mean(loss_metrics['lpips_vgg'])
return loss_metrics
def train(opt):
print("Random Seed: ", opt.seed)
random.seed(opt.seed)
torch.manual_seed(opt.seed)
cudnn.benchmark = True
device = 'cuda'
batch_size = int(opt.batch_size)
domain = opt.domain
# tensorboard
os.makedirs(os.path.join(opt.outf, 'runs'), exist_ok=True)
writer = SummaryWriter(log_dir=os.path.join(opt.outf, 'runs'))
# datasets
train_transform = data.get_transform(domain, 'imtrain')
train_dset = data.get_dataset(domain,
'train',
load_w=True,
transform=train_transform)
print("Training transform:")
print(train_transform)
val_transform = data.get_transform(domain, 'imval')
val_dset = data.get_dataset(domain,
'val',
load_w=True,
transform=val_transform)
print("Validation transform:")
print(val_transform)
train_loader = DataLoader(train_dset,
batch_size=opt.batch_size,
shuffle=True,
pin_memory=False,
num_workers=opt.workers)
val_loader = DataLoader(val_dset,
batch_size=opt.batch_size,
shuffle=False,
pin_memory=False,
num_workers=opt.workers)
# classifier: resnet18 model
net = torchvision.models.resnet18(
num_classes=len(train_dset.coarse_labels))
# load the model weights to finetune from
ckpt_path = ('results/classifiers/%s/%s/net_best.pth' %
(opt.domain, opt.finetune_from))
print("Finetuning model from %s" % ckpt_path)
ckpt = torch.load(ckpt_path)
state_dict = ckpt['state_dict']
net.load_state_dict(state_dict)
net = net.to(device)
# losses + optimizers
criterion = nn.CrossEntropyLoss().to(device) # loss(output, target)
optimizer = optim.Adam(net.parameters(),
lr=opt.lr,
betas=(opt.beta1, 0.999))
scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer,
mode='max',
patience=10,
min_lr=1e-6,
verbose=True)
start_ep = 0
best_val_acc = 0.0
best_val_epoch = 0
# val tensor transform does not flip, train tensor transform
# mimics random resized crop and random flip
val_tensor_transform = data.get_transform(domain, 'tensorbase')
train_tensor_transform = data.get_transform(domain, 'tensormixedtrain')
# load GAN
generator = domain_generator.define_generator('stylegan2', domain)
for epoch in pbar(range(start_ep, opt.niter + 1)):
# average meter for train/val loss and train/val acc
metrics = dict(train_loss=util.AverageMeter(),
val_loss=util.AverageMeter(),
train_acc=util.AverageMeter(),
val_acc=util.AverageMeter())
# train loop
for step, item in enumerate(pbar(train_loader)):
im_orig = item[0].cuda()
opt_w = item[1].cuda()
label = item[2].cuda()
with torch.no_grad():
if opt.perturb_type == 'stylemix':
seed = epoch * len(train_loader) + step
mix_latent = generator.seed2w(n=opt_w.shape[0], seed=seed)
generated_im = generator.perturb_stylemix(
opt_w,
opt.perturb_layer,
mix_latent,
n=opt_w.shape[0],
is_eval=False)
elif opt.perturb_type == 'isotropic':
# used minimum value for isotropic setting
eps = np.min(
generator.perturb_settings['isotropic_eps_%s' %
opt.perturb_layer])
generated_im = generator.perturb_isotropic(
opt_w,
opt.perturb_layer,
eps=eps,
n=opt_w.shape[0],
is_eval=False)
elif opt.perturb_type == 'pca':
# used median value for pca setting
eps = np.median(generator.perturb_settings['pca_eps'])
generated_im = generator.perturb_pca(opt_w,
opt.perturb_layer,
eps=eps,
n=opt_w.shape[0],
is_eval=False)
else:
generated_im = generator.decode(opt_w)
generated_im = train_tensor_transform(generated_im)
# sanity check that the shapes match
assert (generated_im.shape == im_orig.shape)
# with 50% chance, take the original image for
# training the classifier
im = im_orig if torch.rand(1) > 0.5 else generated_im
net.zero_grad()
output = net(im)
loss = criterion(output, label)
accs, _ = accuracy(output, label, topk=(1, ))
metrics['train_loss'].update(loss, n=len(label))
metrics['train_acc'].update(accs[0], n=len(label))
loss.backward()
optimizer.step()
if step % 200 == 0:
pbar.print("%s: %0.6f" %
('train loss', metrics['train_loss'].avg))
pbar.print("%s: %0.6f" %
('train acc', metrics['train_acc'].avg))
# val loop
net = net.eval()
with torch.no_grad():
for step, item in enumerate(pbar(val_loader)):
im_orig = item[0].cuda()
opt_w = item[1].cuda()
label = item[2].cuda()
with torch.no_grad():
# evaluate on the generated image
im = generator.decode(opt_w)
im = val_tensor_transform(im)
assert (im.shape == im_orig.shape)
output = net(im)
loss = criterion(output, label)
accs, _ = accuracy(output, label, topk=(1, ))
metrics['val_loss'].update(loss, n=len(label))
metrics['val_acc'].update(accs[0], n=len(label))
net = net.train()
scheduler.step(metrics['val_acc'].avg)
# send losses to tensorboard
for k, v in metrics.items():
pbar.print("Metrics at end of epoch")
pbar.print("%s: %0.4f" % (k, v.avg))
writer.add_scalar(k.replace('_', '/'), v.avg, epoch)
pbar.print("Learning rate: %0.6f" % optimizer.param_groups[0]['lr'])
# do checkpoint as latest
util.make_checkpoint(net, optimizer, epoch, metrics['val_acc'].avg,
opt.outf, 'latest')
if metrics['val_acc'].avg > best_val_acc:
pbar.print("Updating best checkpoint at epoch %d" % epoch)
pbar.print("Old Best Epoch %d Best Val %0.6f" %
(best_val_epoch, best_val_acc))
# do checkpoint as best
util.make_checkpoint(net, optimizer, epoch, metrics['val_acc'].avg,
opt.outf, 'best')
best_val_acc = metrics['val_acc'].avg
best_val_epoch = epoch
pbar.print("New Best Epoch %d Best Val %0.6f" %
(best_val_epoch, best_val_acc))
with open("%s/best_val.txt" % opt.outf, "w") as f:
f.write("Best Epoch %d Best Val %0.6f\n" %
(best_val_epoch, best_val_acc))
if (float(optimizer.param_groups[0]['lr']) <= 1e-6
and epoch >= best_val_epoch + 50):
pbar.print("Exiting training")
pbar.print("Best Val epoch %d" % best_val_epoch)
pbar.print("Curr epoch %d" % epoch)
break
def train(opt):
print("Random Seed: ", opt.seed)
random.seed(opt.seed)
torch.manual_seed(opt.seed)
cudnn.benchmark = True
device = 'cuda'
batch_size = int(opt.batch_size)
domain = opt.domain
# tensorboard
os.makedirs(os.path.join(opt.outf, 'runs'), exist_ok=True)
writer = SummaryWriter(log_dir=os.path.join(opt.outf, 'runs'))
# datasets
train_transform = data.get_transform(domain, 'imtrain')
train_dset = data.get_dataset(domain,
'train',
load_w=False,
transform=train_transform)
print("Training transform:")
print(train_transform)
val_transform = data.get_transform(domain, 'imval')
val_dset = data.get_dataset(domain,
'val',
load_w=False,
transform=val_transform)
print("Validation transform:")
print(val_transform)
train_loader = DataLoader(train_dset,
batch_size=opt.batch_size,
shuffle=True,
pin_memory=False,
num_workers=opt.workers)
val_loader = DataLoader(val_dset,
batch_size=opt.batch_size,
shuffle=False,
pin_memory=False,
num_workers=opt.workers)
# classifier: resnet18 model
net = torchvision.models.resnet18(
num_classes=len(train_dset.coarse_labels))
if not opt.train_from_scratch:
state_dict = torchvision.models.utils.load_state_dict_from_url(
torchvision.models.resnet.model_urls['resnet18'])
del state_dict['fc.weight']
del state_dict['fc.bias']
net.load_state_dict(state_dict, strict=False)
net = net.to(device)
# losses + optimizers + scheduler
# use smaller learning rate for the feature layers if initialized
# with imagenet pretrained weights
criterion = nn.CrossEntropyLoss().to(device) # loss(output, target)
fc_params = [k[1] for k in net.named_parameters() if k[0].startswith('fc')]
feat_params = [
k[1] for k in net.named_parameters() if not k[0].startswith('fc')
]
feature_backbone_lr = opt.lr if opt.train_from_scratch else 0.1 * opt.lr
print("Initial learning rate for feature backbone: %0.4f" %
feature_backbone_lr)
print("Initial learning rate for FC layer: %0.4f" % opt.lr)
optimizer = optim.Adam([{
'params': fc_params
}, {
'params': feat_params,
'lr': feature_backbone_lr
}],
lr=opt.lr,
betas=(opt.beta1, 0.999))
scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer,
mode='max',
patience=10,
min_lr=1e-6,
verbose=True)
start_ep = 0
best_val_acc = 0.0
best_val_epoch = 0
for epoch in pbar(range(start_ep, opt.niter + 1)):
# average meter for train/val loss and train/val acc
metrics = dict(train_loss=util.AverageMeter(),
val_loss=util.AverageMeter(),
train_acc=util.AverageMeter(),
val_acc=util.AverageMeter())
# train loop
for step, item in enumerate(pbar(train_loader)):
im = item[0].cuda()
label = item[1].cuda()
net.zero_grad()
output = net(im)
loss = criterion(output, label)
accs, _ = accuracy(output, label, topk=(1, ))
metrics['train_loss'].update(loss, n=len(label))
metrics['train_acc'].update(accs[0], n=len(label))
loss.backward()
optimizer.step()
pbar.print("%s: %0.2f" % ('train loss', metrics['train_loss'].avg))
pbar.print("%s: %0.2f" % ('train acc', metrics['train_acc'].avg))
# val loop
net = net.eval()
with torch.no_grad():
for step, item in enumerate(pbar(val_loader)):
im = item[0].cuda()
label = item[1].cuda()
output = net(im)
loss = criterion(output, label)
accs, _ = accuracy(output, label, topk=(1, ))
metrics['val_loss'].update(loss, n=len(label))
metrics['val_acc'].update(accs[0], n=len(label))
net = net.train()
# update scheduler
scheduler.step(metrics['val_acc'].avg)
# send losses to tensorboard
for k, v in metrics.items():
pbar.print("Metrics at end of epoch")
pbar.print("%s: %0.4f" % (k, v.avg))
writer.add_scalar(k.replace('_', '/'), v.avg, epoch)
pbar.print("Learning rate: %0.6f" % optimizer.param_groups[0]['lr'])
# do checkpoint as latest
util.make_checkpoint(net, optimizer, epoch, metrics['val_acc'].avg,
opt.outf, 'latest')
if metrics['val_acc'].avg > best_val_acc:
pbar.print("Updating best checkpoint at epoch %d" % epoch)
pbar.print("Old Best Epoch %d Best Val %0.2f" %
(best_val_epoch, best_val_acc))
# do checkpoint as best
util.make_checkpoint(net, optimizer, epoch, metrics['val_acc'].avg,
opt.outf, 'best')
best_val_acc = metrics['val_acc'].avg
best_val_epoch = epoch
pbar.print("New Best Epoch %d Best Val %0.2f" %
(best_val_epoch, best_val_acc))
with open("%s/best_val.txt" % opt.outf, "w") as f:
f.write("Best Epoch %d Best Val %0.2f\n" %
(best_val_epoch, best_val_acc))
# terminate training if reached min LR and best validation is
# not improving
if (float(optimizer.param_groups[0]['lr']) <= 1e-6
and epoch >= best_val_epoch + 20):
pbar.print("Exiting training")
pbar.print("Best Val epoch %d" % best_val_epoch)
pbar.print("Curr epoch %d" % epoch)
break
def train(self, dataset):
g_train_loss = metrics.Mean()
d_train_loss = metrics.Mean()
current_time = datetime.now().strftime("%Y%m%d-%H%M%S")
checkpoint_directory = "./checkpoints/training_checkpoints"
g_checkpoint_prefix = os.path.join(checkpoint_directory + "/generator",
"ckpt")
d_checkpoint_prefix = os.path.join(
checkpoint_directory + "/discriminator", "ckpt")
train_log_dir = 'logs/gradient_tape/' + current_time + '/train'
train_summary_writer = tf.summary.create_file_writer(train_log_dir)
for epoch in tf.range(self.epochs):
epoch = tf.cast(epoch, dtype=tf.int64, name=epoch)
bar = pbar(self.total_passwords, self.batch_size, epoch,
self.epochs)
for iteration, batch in zip(range(self.iterations), dataset):
for _ in tf.range(self.n_critic):
self.text = batch['password']
real = tf.reshape(tf.dtypes.cast(self.text, tf.float32),
[2, 1, 32])
self.train_d(real)
d_loss = self.train_d(real)
d_train_loss(d_loss)
g_loss = self.train_g()
g_train_loss(g_loss)
self.train_g()
"""
Tensorboard tracking calls
files generated sent to
/logs/gradient_tape/
"""
with train_summary_writer.as_default():
tf.summary.scalar('Generator',
g_train_loss.result(),
step=epoch)
tf.summary.scalar('Accuracy',
d_train_loss.result(),
step=epoch)
bar.postfix['g_loss'] = f'{g_train_loss.result():6.3f}'
bar.postfix['d_loss'] = f'{d_train_loss.result():6.3f}'
bar.update(self.batch_size)
if iteration % self.checkpoints == 0 and iteration > 0:
generator_checkpoint = tf.train.Checkpoint(
optimizer=self.g_opt, model=self.G)
generator_checkpoint.save(file_prefix=g_checkpoint_prefix)
discriminator_checkpoint = tf.train.Checkpoint(
optimizer=self.d_opt, model=self.D)
discriminator_checkpoint.save(
file_prefix=d_checkpoint_prefix)
self.G.summary()
self.D.summary()
"""
Tensorflow model save
files located at: /models/generator or /models/discriminator
"""
tf.saved_model.save(
self.G,
'./models/generator/' + self.dataset_name + current_time)
tf.saved_model.save(
self.D,
'./models/discriminator/' + self.dataset_name + current_time)
g_train_loss.reset_states()
d_train_loss.reset_states()
bar.close()
del bar