def extract_statistics(cfg, train_set, inliner_classes, E, G):
zlist = []
rlist = []
data_loader = make_dataloader(train_set, cfg.TEST.BATCH_SIZE, torch.cuda.current_device())
for label, x in data_loader:
x = x.view(-1, cfg.MODEL.INPUT_IMAGE_SIZE * cfg.MODEL.INPUT_IMAGE_SIZE)
z,logvar,mu = E(x.view(-1, 1, cfg.MODEL.INPUT_IMAGE_SIZE, cfg.MODEL.INPUT_IMAGE_SIZE))
recon_batch = G(z)
z = z.squeeze()
recon_batch = recon_batch.squeeze().cpu().detach().numpy()
x = x.squeeze().cpu().detach().numpy()
z = z.cpu().detach().numpy()
for i in range(x.shape[0]):
distance = np.linalg.norm(x[i].flatten() - recon_batch[i].flatten())
rlist.append(distance)
zlist.append(z)
zlist = np.concatenate(zlist)
counts, bin_edges = np.histogram(rlist, bins=30, normed=True)
if cfg.MAKE_PLOTS:
plt.plot(bin_edges[1:], counts, linewidth=2)
save_plot(r"Distance, $\left \|\| I - \hat{I} \right \|\|$",
'Probability density',
r"PDF of distance for reconstruction error, $p\left(\left \|\| I - \hat{I} \right \|\| \right)$",
'mnist_%s_reconstruction_error.pdf' % ("_".join([str(x) for x in inliner_classes])))
for i in range(cfg.MODEL.LATENT_SIZE):
plt.hist(zlist[:, i], bins='auto', histtype='step')
if cfg.MAKE_PLOTS:
save_plot(r"$z$",
'Probability density',
r"PDF of embeding $p\left(z \right)$",
'mnist_%s_embedding.pdf' % ("_".join([str(x) for x in inliner_classes])))
def fmin(func, x0, args, disp):
x0 = [2.0, 0.0, 1.0]
return scipy.optimize.fmin(func, x0, args, xtol=1e-12, ftol=1e-12, disp=0)
gennorm_param = np.zeros([3, cfg.MODEL.LATENT_SIZE])
for i in range(cfg.MODEL.LATENT_SIZE):
betta, loc, scale = scipy.stats.gennorm.fit(zlist[:, i], optimizer=fmin)
gennorm_param[0, i] = betta
gennorm_param[1, i] = loc
gennorm_param[2, i] = scale
return counts, bin_edges, gennorm_param
def run_novely_prediction_on_dataset(dataset,
percentage,
concervative=False):
dataset.shuffle()
dataset = create_set_with_outlier_percentage(dataset, inliner_classes,
percentage, concervative)
result = []
gt_novel = []
data_loader = make_dataloader(dataset, cfg.TEST.BATCH_SIZE,
torch.cuda.current_device())
include_jacobian = True
N = (cfg.MODEL.INPUT_IMAGE_SIZE * cfg.MODEL.INPUT_IMAGE_SIZE -
cfg.MODEL.LATENT_SIZE) * mul
logC = loggamma(N / 2.0) - (N / 2.0) * np.log(2.0 * np.pi)
def logPe_func(x):
# p_{\|W^{\perp}\|} (\|w^{\perp}\|)
# \| w^{\perp} \|}^{m-n}
return logC - (N - 1) * np.log(x) + np.log(
r_pdf(x, bin_edges, counts))
for label, x in data_loader:
x = x.view(
-1, cfg.MODEL.INPUT_IMAGE_CHANNELS *
cfg.MODEL.INPUT_IMAGE_SIZE * cfg.MODEL.INPUT_IMAGE_SIZE)
x = Variable(x.data, requires_grad=True)
z = E(
x.view(-1, cfg.MODEL.INPUT_IMAGE_CHANNELS,
cfg.MODEL.INPUT_IMAGE_SIZE, cfg.MODEL.INPUT_IMAGE_SIZE))
recon_batch = G(z)
z = z.squeeze()
if include_jacobian:
J = compute_jacobian_autograd(x, z)
J = J.cpu().numpy()
z = z.cpu().detach().numpy()
recon_batch = recon_batch.squeeze().cpu().detach().numpy()
x = x.squeeze().cpu().detach().numpy()
for i in range(x.shape[0]):
if include_jacobian:
u, s, vh = np.linalg.svd(J[i, :, :], full_matrices=False)
logD = -np.sum(np.log(
np.abs(s))) # | \mathrm{det} S^{-1} |
# logD = np.log(np.abs(1.0/(np.prod(s))))
else:
logD = 0
p = scipy.stats.gennorm.pdf(z[i], gennorm_param[0, :],
gennorm_param[1, :],
gennorm_param[2, :])
logPz = np.sum(np.log(p))
# Sometimes, due to rounding some element in p may be zero resulting in Inf in logPz
# In this case, just assign some large negative value to make sure that the sample
# is classified as unknown.
if not np.isfinite(logPz):
logPz = -1000
distance = np.linalg.norm(x[i].flatten() -
recon_batch[i].flatten())
logPe = logPe_func(distance)
P = logD + logPz + logPe
result.append(P)
gt_novel.append(label[i].item() in inliner_classes)
result = np.asarray(result, dtype=np.float32)
ground_truth = np.asarray(gt_novel, dtype=np.float32)
return result, ground_truth
def train(folding_id, inliner_classes, ic):
cfg = get_cfg_defaults()
cfg.merge_from_file('configs/mnist.yaml')
cfg.freeze()
logger = logging.getLogger("logger")
zsize = cfg.MODEL.LATENT_SIZE
output_folder = os.path.join('results_' + str(folding_id) + "_" +
"_".join([str(x) for x in inliner_classes]))
os.makedirs(output_folder, exist_ok=True)
os.makedirs('models', exist_ok=True)
train_set, _, _ = make_datasets(cfg, folding_id, inliner_classes)
logger.info("Train set size: %d" % len(train_set))
G = Generator(cfg.MODEL.LATENT_SIZE,
channels=cfg.MODEL.INPUT_IMAGE_CHANNELS)
G.weight_init(mean=0, std=0.02)
D = Discriminator(channels=cfg.MODEL.INPUT_IMAGE_CHANNELS)
D.weight_init(mean=0, std=0.02)
E = Encoder(cfg.MODEL.LATENT_SIZE, channels=cfg.MODEL.INPUT_IMAGE_CHANNELS)
E.weight_init(mean=0, std=0.02)
if cfg.MODEL.Z_DISCRIMINATOR_CROSS_BATCH:
ZD = ZDiscriminator_mergebatch(zsize, cfg.TRAIN.BATCH_SIZE)
else:
ZD = ZDiscriminator(zsize, cfg.TRAIN.BATCH_SIZE)
ZD.weight_init(mean=0, std=0.02)
lr = cfg.TRAIN.BASE_LEARNING_RATE
G_optimizer = optim.Adam(G.parameters(), lr=lr, betas=(0.5, 0.999))
D_optimizer = optim.Adam(D.parameters(), lr=lr, betas=(0.5, 0.999))
GE_optimizer = optim.Adam(list(E.parameters()) + list(G.parameters()),
lr=lr,
betas=(0.5, 0.999))
ZD_optimizer = optim.Adam(ZD.parameters(), lr=lr, betas=(0.5, 0.999))
BCE_loss = nn.BCELoss()
sample = torch.randn(64, zsize).view(-1, zsize, 1, 1)
tracker = LossTracker(output_folder=output_folder)
for epoch in range(cfg.TRAIN.EPOCH_COUNT):
G.train()
D.train()
E.train()
ZD.train()
epoch_start_time = time.time()
data_loader = make_dataloader(train_set, cfg.TRAIN.BATCH_SIZE,
torch.cuda.current_device())
train_set.shuffle()
if (epoch + 1) % 30 == 0:
G_optimizer.param_groups[0]['lr'] /= 4
D_optimizer.param_groups[0]['lr'] /= 4
GE_optimizer.param_groups[0]['lr'] /= 4
ZD_optimizer.param_groups[0]['lr'] /= 4
print("learning rate change!")
for y, x in data_loader:
x = x.view(-1, cfg.MODEL.INPUT_IMAGE_CHANNELS,
cfg.MODEL.INPUT_IMAGE_SIZE, cfg.MODEL.INPUT_IMAGE_SIZE)
y_real_ = torch.ones(x.shape[0])
y_fake_ = torch.zeros(x.shape[0])
y_real_z = torch.ones(
1 if cfg.MODEL.Z_DISCRIMINATOR_CROSS_BATCH else x.shape[0])
y_fake_z = torch.zeros(
1 if cfg.MODEL.Z_DISCRIMINATOR_CROSS_BATCH else x.shape[0])
#############################################
D.zero_grad()
D_result = D(x).squeeze()
D_real_loss = BCE_loss(D_result, y_real_)
z = torch.randn((x.shape[0], zsize)).view(-1, zsize, 1, 1)
z = Variable(z)
x_fake = G(z).detach()
D_result = D(x_fake).squeeze()
D_fake_loss = BCE_loss(D_result, y_fake_)
D_train_loss = D_real_loss + D_fake_loss
D_train_loss.backward()
D_optimizer.step()
tracker.update(dict(D=D_train_loss))
#############################################
G.zero_grad()
z = torch.randn((x.shape[0], zsize)).view(-1, zsize, 1, 1)
z = Variable(z)
x_fake = G(z)
D_result = D(x_fake).squeeze()
G_train_loss = BCE_loss(D_result, y_real_)
G_train_loss.backward()
G_optimizer.step()
tracker.update(dict(G=G_train_loss))
#############################################
ZD.zero_grad()
z = torch.randn((x.shape[0], zsize)).view(-1, zsize)
z = Variable(z)
ZD_result = ZD(z).squeeze()
ZD_real_loss = BCE_loss(ZD_result, y_real_z)
z = E(x).squeeze().detach()
ZD_result = ZD(z).squeeze()
ZD_fake_loss = BCE_loss(ZD_result, y_fake_z)
ZD_train_loss = ZD_real_loss + ZD_fake_loss
ZD_train_loss.backward()
ZD_optimizer.step()
tracker.update(dict(ZD=ZD_train_loss))
# #############################################
E.zero_grad()
G.zero_grad()
z = E(x)
x_d = G(z)
ZD_result = ZD(z.squeeze()).squeeze()
E_train_loss = BCE_loss(ZD_result, y_real_z) * 1.0
Recon_loss = F.binary_cross_entropy(x_d, x.detach()) * 2.0
(Recon_loss + E_train_loss).backward()
GE_optimizer.step()
tracker.update(dict(GE=Recon_loss, E=E_train_loss))
# #############################################
comparison = torch.cat([x, x_d])
save_image(comparison.cpu(),
os.path.join(output_folder,
'reconstruction_' + str(epoch) + '.png'),
nrow=x.shape[0])
epoch_end_time = time.time()
per_epoch_ptime = epoch_end_time - epoch_start_time
logger.info(
'[%d/%d] - ptime: %.2f, %s' %
((epoch + 1), cfg.TRAIN.EPOCH_COUNT, per_epoch_ptime, tracker))
tracker.register_means(epoch)
tracker.plot()
with torch.no_grad():
resultsample = G(sample).cpu()
save_image(
resultsample.view(64, cfg.MODEL.INPUT_IMAGE_CHANNELS,
cfg.MODEL.INPUT_IMAGE_SIZE,
cfg.MODEL.INPUT_IMAGE_SIZE),
os.path.join(output_folder, 'sample_' + str(epoch) + '.png'))
logger.info("Training finish!... save training results")
os.makedirs("models", exist_ok=True)
print("Training finish!... save training results")
torch.save(G.state_dict(), "models/Gmodel_%d_%d.pkl" % (folding_id, ic))
torch.save(E.state_dict(), "models/Emodel_%d_%d.pkl" % (folding_id, ic))