def estimate_log_likelihood(data_loader, device, split="valid"):
model_path = "best_model.pth"
model = VAE().to(device)
model.load_state_dict(torch.load(model_path, map_location=device))
# The number of importance samples
K = 200
model.eval()
log_likelihood_estimate = 0.0
with torch.no_grad():
for batch_idx, mini_batch_x in enumerate(data_loader):
mini_batch_x = mini_batch_x.to(device)
recon_output, mu, logvar = model(mini_batch_x)
# sample K=200 importance samples from posterior q(z|x_i)
# Z is of size (M, K, L)
Z = generate_K_samples(mu, logvar, K)
# estimate log-likelihood for a batch
log_px = importance_sampling(model, mini_batch_x, Z)
log_likelihood_estimate += log_px.sum()
print('log_px estimate of mini-batch {} of {} set: {:.4f}'.format(
batch_idx, split, log_px.mean()))
print('log_px estimate of {} set: {:.4f}'.format(
split, log_likelihood_estimate / len(data_loader.dataset)))
def get_model(path):
ckpt = torch.load(path)
train_args = ckpt['args']
# model = {'dae': DAE, 'vae': VAE, 'aae': AAE}[train_args.model](
# vocab, train_args).to(device)
model = VAE(vocab, train_args, device)
model.load_state_dict(ckpt['model'])
model.flatten()
model.eval()
return model
def save_z(path):
data = AwA2()
loader = DataLoader(dataset=data, batch_size=batch_size, shuffle=False)
net = VAE(z_dim)
net = net.to(device)
net.load_state_dict(torch.load(path))
print('Load model from%s' % path)
net.eval()
z_list = np.empty([0, z_dim])
for f, label in loader:
print(label)
f = f.to(device)
f = f.float()
recon_f, mu, logvar = net(f)
z = reparametrize(mu, logvar)
z = z.detach().cpu().numpy()
z_list = np.append(z_list, z, axis=0)
print(z_list.shape)
np.save('Data/b%.1fVAE-%d.npy' % (beta, z_dim), z_list)
def generate(model_path, device):
print("Loading model....\n")
model = VAE().to(device)
model.load_state_dict(torch.load(model_path, map_location=device))
model.eval()
with torch.no_grad():
z = torch.randn(10, 100).to(device)
decoding = model.fc_decode(z)
decoding = decoding.reshape(decoding.shape[0], decoding.shape[1], 1, 1)
gen_output = model.decoder(decoding)
gen_output = gen_output.squeeze().detach().cpu().numpy()
gen_output[gen_output >= 0.5] = 1
gen_output[gen_output < 0.5] = 0
# plotting
fig, axs = plt.subplots(2, 5)
fig.suptitle('Generated images')
for i in range(2):
for j in range(5):
axs[i, j].imshow(gen_output[j + i * 5])
plt.savefig("Generated_Samples.png")
def importance_sampling(VAE, inputs, K=200):
sum_arg = []
VAE.eval()
with torch.no_grad():
inputs = inputs.to(DEVICE)
_, mean, log_sigma = VAE.encoder.forward(inputs)
sigma = torch.exp(log_sigma)
z = torch.randn(
(inputs.shape[0], K, mean.shape[1])).to(DEVICE) #shape(M, K, L)
for samp in range(K):
z_i = mean + sigma * z[:, samp, :]
p_vals = gaussian_density(z_i, 0, 1)
q_vals = gaussian_density(z_i, mean, sigma)
recons = VAE.decoder.forward(z_i)
log_p_z = torch.sum(torch.log(p_vals), dim=1)
log_q_z = torch.sum(torch.log(q_vals), dim=1)
log_p_xz = reconstruction_loss(recons, inputs).squeeze(1)
sumarg = log_p_xz + log_p_z - log_q_z
sumarg = sumarg.cpu().numpy()
sum_arg.append(sumarg)
return -np.log(K) + logsumexp(sum_arg, axis=0)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
weighted_loss = 1
ip = "127.0.0.1" # Ip address that the TCP/IP interface listens to
port = 13000 # Port number that the TCP/IP interface listens to
size = 64 # Please check the Updates section above for more details
timescale = 10 # Please check the Updates section above for more details
env = Neurosmash.Environment(timescale=timescale, size=size, port=port, ip=ip)
# Load VAE weights
vae = VAE(device, image_channels=3).to(device)
vae.load_state_dict(
torch.load("./data_folder_vae/vae_v3_weighted_loss_{}.torch".format(
weighted_loss)))
vae.eval()
# Load RNN weights
rnn = MDNRNN(32, 256, 5, 1).to(device)
rnn.load_state_dict(
torch.load("./weights/rnn_29dec_{}.torch".format(weighted_loss)))
rnn.eval()
# Load controller, if vanilla DQN, replace DQN_VAE with DQN2
if USE_WM:
if USE_RNN:
input_params = n_actions * 256 + 32
else:
input_params = 32
policy_net = DQN_VAE(64, 64, 3, input_params).to(device)
#Initialize model
if m == 'VAE':
model = VAE(hyperparams)
if m == 'IWAE':
model = IWAE(hyperparams)
#Evaluate
print 'Evaluating'
start = time.time()
# test_results, train_results, test_labels, train_labels = model.eval(data=test_x, batch_size=n_batch_eval, display_step=100,
# path_to_load_variables=parameter_path+saved_parameter_file, data2=train_x)
test_results, train_results = model.eval(
data=test_x,
batch_size=n_batch_eval,
display_step=100,
path_to_load_variables=parameter_path +
saved_parameter_file,
data2=train_x,
k_eval=k_evaluation)
time_to_eval = time.time() - start
print 'time to evaluate', time_to_eval
print 'Results: ' + str(test_results) + 'train:' + str(
train_results) + ' for ' + saved_parameter_file
if save_log:
with open(experiment_log, "a") as myfile:
myfile.write('time to evaluate ' +
str(time_to_eval) + '\n')
myfile.write('test set LL ' + str(test_results) +
transform = transforms.Compose([transforms.ToTensor()])
# load the last checkpoint, if it exists
if not exists(join(args.model_dir, args.ckpt)):
raise Exception("model does not exist")
location = 'cuda'
checkpoint = torch.load(join(args.model_dir, args.ckpt), map_location=location)
model.load_state_dict(checkpoint['model_state_dict'])
# optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
# validation_iwae = checkpoint['validation_iwae']
# train_vlb = checkpoint['train_vlb']
model = model.to(device)
model.eval()
print("model loaded")
# size = 10.0
mask = 30
# mask size used during latent optimization
for size in [10.0]:
for intensity in [0.4]:
for mask_val in ["zero"]:
for test_mask in [mask]: # [25, 50, 100, 200, 512]:
for fe in ["vgg"]:
for mse in [0.5]:
for run in [9]:
healthy_query = FairDataset(
class NetworkTrainer(object):
def __init__(self, args):
self.GPU_ID = args.gpu_id
self.IS_TRAIN = args.is_train
self.NET = args.net_type
self.LEARNING_RATE_TYPE = args.learning_rate
self.BATCH_SIZE = args.batch_size
self.NUM_EPOCHS = args.num_epochs
self.NUM_CALC_EPOCHS = args.num_calc_epochs
self.SAVE_SET = args.save_set
self.CALC_EPOCHS = torch.linspace(0, self.NUM_EPOCHS,
self.NUM_CALC_EPOCHS + 1)
self.DATA_ROOT = DATA_ROOT
self.DATA_SET = DATA_SET
self.LATENT_DIM = LATENT_DIM
self.LR_LIST_3 = np.ones(args.num_epochs) * 1e-3 # arg:'3'
self.LR_LIST_45 = np.logspace(-4, -5, args.num_epochs) # arg:'45'
self.PATH = {}
self.PATH['fig_save_path'] = PROJECT_ROOT + '/save/figs/'
self.PATH[
'fig_save_path_curve'] = self.PATH['fig_save_path'] + 'curve.png'
self.PATH['model_save_path'] = PROJECT_ROOT + '/save/models/'
self.PATH[
'model_save_path_net'] = self.PATH['model_save_path'] + 'model.pkl'
self.PATH[
'data_save_path'] = PROJECT_ROOT + '/save/data/{}_{}/'.format(
self.NET, self.DATA_SET)
self.PATH[
'data_save_path_log'] = self.PATH['data_save_path'] + 'log.pkl'
if not os.path.exists(self.PATH['fig_save_path']):
os.makedirs(self.PATH['fig_save_path'])
if not os.path.exists(self.PATH['model_save_path']):
os.makedirs(self.PATH['model_save_path'])
if not os.path.exists(self.PATH['data_save_path']):
os.makedirs(self.PATH['data_save_path'])
self.EPOCH = 1
def prepare(self):
torch.cuda.set_device(self.GPU_ID)
self._prepare_model()
self._prepare_dataset()
self._generate_plot_dic()
def _prepare_model(self):
dim = np.load(os.path.join(self.DATA_ROOT, "train", "dim.npy"))
if self.NET == 'VAE':
self.model = VAE(input_dim=dim,
latent_dim=self.LATENT_DIM).to(self.GPU_ID)
if self.LEARNING_RATE_TYPE == '45':
self.LEARNING_RATE = self.LR_LIST_45
elif self.LEARNING_RATE_TYPE == '3':
self.LEARNING_RATE = self.LR_LIST_3
self.optimizer = optim.Adam(self.model.parameters(),
lr=self.LEARNING_RATE[0],
betas=(0.9, 0.999),
eps=1e-08,
weight_decay=0)
def _prepare_dataset(self):
print(self.DATA_ROOT)
# load the data
train_dataset = BinaryDataset(self.DATA_ROOT, phase='train')
test_dataset = BinaryDataset(self.DATA_ROOT, phase='test')
# encapsulate them into dataloader form
self.train_loader = torch.utils.data.DataLoader(
train_dataset, batch_size=self.BATCH_SIZE, drop_last=True)
self.test_loader = torch.utils.data.DataLoader(
test_dataset,
batch_size=self.BATCH_SIZE,
shuffle=False,
drop_last=True)
def _generate_plot_dic(self):
self.plot_dict = {'train_loss': [], 'test_loss': []}
def train_epoch(self):
self.model.train()
print("Learning rate is", self.LEARNING_RATE[self.EPOCH - 1])
update_lr(self.optimizer, self.LEARNING_RATE[self.EPOCH - 1])
total_loss = 0
cnt = 0
for record in tqdm(self.train_loader,
desc="epoch " + str(self.EPOCH),
mininterval=1):
record = Variable(record).to(self.GPU_ID)
self.optimizer.zero_grad()
if self.EPOCH == 1 and cnt == 0:
self.save(epoch=0)
cnt += 1
recon_x, mu, logvar = self.model.forward(record)
loss = loss_func(recon_x, record, mu, logvar)
loss.backward()
total_loss += loss
self.optimizer.step()
length = len(self.train_loader) // self.BATCH_SIZE
avg_loss = float(total_loss) / length
self.plot_dict['train_loss'].append(avg_loss)
print(
'LOSS of the network on the training records: {}'.format(avg_loss))
return avg_loss
def eval_epoch(self):
self.model.eval()
total_loss = 0
with torch.no_grad():
for record in tqdm(self.test_loader,
desc="evaluation " + str(self.EPOCH),
mininterval=1):
record = Variable(record).to(self.GPU_ID)
recon_x, mu, logvar = self.model.forward(record)
loss = loss_func(recon_x, record, mu, logvar)
total_loss += loss
length = len(self.test_loader.dataset) // self.BATCH_SIZE
avg_loss = float(total_loss) / length
self.plot_dict['test_loss'].append(avg_loss)
print(
'LOSS of the network on the testing records: {}'.format(avg_loss))
return avg_loss
def draw(self):
print('ploting...')
plt.figure(figsize=(16, 8))
plt.subplot(1, 2, 1)
x = range(1, len(self.plot_dict["train_loss"]) + 1)
plt.xlabel("epoch")
plt.plot(x, self.plot_dict["train_loss"], label="train_loss")
plt.legend()
plt.subplot(1, 2, 2)
x = range(1, len(self.plot_dict["test_loss"]) + 1)
plt.xlabel("epoch")
plt.plot(x, self.plot_dict["test_loss"], label="test_loss")
plt.legend()
plt.tight_layout()
plt.savefig(self.PATH['fig_save_path_curve'],
bbox_inches='tight',
dpi=300)
def save(self, epoch):
print('saving...')
torch.save(self.model.state_dict(),
self.PATH['model_save_path'] + 'model_{}.pkl'.format(epoch))
with open(self.PATH['data_save_path_log'], 'wb') as f:
pickle.dump(self.plot_dict, f)
def run(self):
self.prepare()
for self.EPOCH in range(1, self.NUM_EPOCHS + 1):
self.train_epoch()
self.eval_epoch()
if self.EPOCH % 10 == 0:
self.draw()
self.save(self.EPOCH)
class Donut:
def __init__(self, CKPT_path=None):
self._vae = VAE()
self.optimizer = Adam(self._vae.parameters(),
lr=0.001,
weight_decay=0.001)
if CKPT_path is not None:
model_CKPT = torch.load(CKPT_path)
self._vae.load_state_dict(model_CKPT['state_dict'])
self.optimizer.load_state_dict(model_CKPT['optimizer'])
print('load vae and optimizer form file')
def m_elbo_loss(self, train_x, train_y, z, x_miu, x_std, z_miu, z_std):
"""
采用蒙特卡洛估计， 根据论文，设置采样次数 L=1
:param train_x: batch_size * win, 样本
:param train_y: batch_size * 1, 标签， 0代表正常， 1代表异常
:param z: batch_size * latent_size, 观察到的隐变量
:param x_miu: batch_size * win, x服从正态分布的均值
:param x_std: batch_size * win, x服从正态分布的标准差
:param z_miu: batch_size * latent_size, 后验z服从正态分布的均值
:param z_std: batch_size * latent_size, 后验z服从正态分布的标准差
:param z_prior_mean: int, 先验z服从正态分布的均值， 一般设置为0
:param z_prior_std: int, 先验z服从正态分布的标准差， 一般设置为1
:return:
"""
z_prior_mean = torch.zeros(size=z_miu.shape)
z_prior_std = torch.ones(size=z_miu.shape)
# 以下蒙特卡洛估计的采样次数均为1
# 蒙特卡洛估计 log p(x|z)。在重构的x的正态分布上， 取值为train_x时，概率密度函数的log值； batch_size * win
log_p_xz = -torch.log(math.sqrt(2 * math.pi) * x_std) - (
(train_x - x_miu)**2) / (2 * x_std**2)
# 蒙特卡洛估计 log p(z). p(z)为先验分布， 一般设置为标准正态分布； batch_size * latent_size
log_p_z = -torch.log(math.sqrt(2 * math.pi) * z_prior_std) - (
(z - z_prior_mean)**2) / (2 * z_prior_std**2)
# 蒙特卡洛估计 log q(z|x). q(z|x)为z的后验分布; batch_size * latent_size
log_q_zx = -torch.log(math.sqrt(2 * math.pi) * z_std) - (
(z - z_miu)**2) / (2 * z_std**2)
# 去除缺失点的影响，也是m-elbo的精髓
normal = 1 - train_y # batch_size * win
log_p_xz = normal * log_p_xz # batch_size * win
# beta, log_p_z 的放缩系数
beta = torch.sum(normal, dim=1) / normal.shape[1] # size = batch_size
# m-elbo的值
m_elbo = torch.sum(
log_p_xz, dim=1) + beta * torch.sum(log_p_z, dim=1) - torch.sum(
log_q_zx, dim=1)
m_elbo = torch.mean(m_elbo) * (-1)
return m_elbo
def fit(self, x, y, n_epoch=30, valid_x=None, valid_y=None):
'''
如果在实际应用中没有标签， 可以把大多数样本当做正常样本， 即把y全部设置为0
:param x:
:param y:
:param n_epoch:
:param valid_x:
:param valid_y:
:return:
'''
# todo missing injection
self._vae.train()
train_dataset = TsDataset(x, y)
train_iter = torch.utils.data.DataLoader(train_dataset,
batch_size=256,
shuffle=True,
num_workers=0)
if valid_x is not None:
valid_dataset = TsDataset(valid_x, valid_y)
valid_iter = torch.utils.data.DataLoader(valid_dataset,
batch_size=128,
shuffle=False,
num_workers=0)
optimizer = self.optimizer # todo 动态学习率
lr_scheduler = StepLR(optimizer, step_size=100, gamma=0.75)
for epoch in range(n_epoch):
lr_scheduler.step()
for train_x, train_y in train_iter:
optimizer.zero_grad()
z, x_miu, x_std, z_miu, z_std = self._vae(train_x) # 前向传播
l = self.m_elbo_loss(train_x, train_y, z, x_miu, x_std, z_miu,
z_std)
l.backward()
optimizer.step()
# 保存模型
if epoch % 100 == 0:
print("保存模型")
torch.save(
{
"state_dict": self._vae.state_dict(),
"optimizer": optimizer.state_dict(),
"loss": l.item()
}, "./model_parameters/epoch{}-loss{:.2f}.tar".format(
epoch, l.item()))
# 验证集
if epoch % 50 == 0 and valid_x is not None:
with torch.no_grad():
flag = 1
for v_x, v_y in valid_iter:
z, x_miu, x_std, z_miu, z_std = self._vae(v_x)
v_l = self.m_elbo_loss(v_x, v_y, z, x_miu, x_std,
z_miu, z_std)
if flag == 1:
flag = 0
v_x_ = v_x[0].view(1, 120)
z, x_miu, x_std, z_miu, z_std = self._vae(v_x_)
restruct_compare_plot(v_x_.view(120),
x_miu.view(120))
print("train loss %.4f, valid loss %.4f" %
(l.item(), v_l.item()))
with open("log.txt", "a") as f:
f.writelines("%d %.4f %.4f\n" %
(epoch, l.item(), v_l.item()))
else:
print("loss", l.item(), " lr,", lr_scheduler.get_last_lr())
def evaluate(self, test_x, test_y):
# todo mcmc
self._vae.eval()
test_dataset = TsDataset(test_x, test_y)
test_iter = torch.utils.data.DataLoader(test_dataset,
batch_size=2560,
shuffle=False,
num_workers=0)
scores = []
with torch.no_grad():
for x, y in test_iter:
# z的采样次数
z, x_miu, x_std, z_miu, z_std = self._vae(x,
n_sample=20) # 前向传播
# 蒙特卡洛估计 E log p(x|z), 重构概率
log_p_xz = -torch.log(math.sqrt(2 * math.pi) * x_std) - (
(x - x_miu)**2) / (2 * x_std**2)
anomaly_score = -torch.mean(log_p_xz[:, :, -1],
dim=0) # 异常分数越大，越可能为异常。
scores.append(anomaly_score)
scores = torch.cat(scores)
# scores = torch.cat((torch.ones(self._vae.win - 1) * torch.min(scores), scores), dim=0)
# todo 使用segment 评判， 需要将时序恢复成原来的样本，而不是滑动窗口取到的片段
assert len(scores) == len(test_dataset)
# 至此， 得到了异常分数和标签， 需要选定一个最好的异常阈值， 选用标准是 best f1 score
best_threshold(np.array(test_y[:, -1]), scores.detach().numpy())
for x, y in test_iter:
x_ = x[0].view(1, 120)
z, x_miu, x_std, z_miu, z_std = self._vae(x_) # 前向传播
restruct_compare_plot(x_.view(120), x_miu.view(120))
break
def main():
china = []
germany = []
japan = []
akamai_ghost = []
apache = []
router = []
trusted = []
untrusted = []
with open('Data/bin/label_location.pkl', 'rb') as f:
location_dict = pickle.load(f)
with open('Data/bin/label_server.pkl', 'rb') as f:
server_dict = pickle.load(f)
with open('Data/bin/label_certificate.pkl', 'rb') as f:
certificate_dict = pickle.load(f)
dim = np.load(os.path.join(DATA_ROOT, "test", "dim.npy"))
pretrained_model = VAE(input_dim = dim, latent_dim = LATENT_DIM).to(GPU_ID)
pretrained_model.load_state_dict(torch.load(MODEL_PATH))
train_dataset = BinaryDataset(DATA_ROOT, phase = 'train')
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size = BATCH_SIZE, shuffle = False, drop_last = True)
pretrained_model.eval()
with torch.no_grad():
cnt = 0
for record in tqdm(train_loader, desc = "Analyzing test set"):
record = Variable(record).to(GPU_ID)
_, mu, _ = pretrained_model.forward(record)
mu = mu.cpu().numpy()
mu = np.squeeze(mu)
if cnt in location_dict["China"]:
china.append(mu)
elif cnt in location_dict["Germany"]:
germany.append(mu)
elif cnt in location_dict["Japan"]:
japan.append(mu)
if cnt in server_dict["AkamaiGHost"]:
akamai_ghost.append(mu)
elif cnt in server_dict["Apache"]:
apache.append(mu)
elif cnt in server_dict["Router"]:
router.append(mu)
if cnt in certificate_dict[True]:
trusted.append(mu)
elif cnt in certificate_dict[False]:
untrusted.append(mu)
cnt = cnt + 1
print("start PCA....")
location = china + germany + japan
location_2d = pca(location)
china = location_2d[:len(china)]
germany = location_2d[len(china):len(china)+len(germany)]
japan = location_2d[len(china)+len(germany):]
server = akamai_ghost + apache + router
server_2d = pca(server)
akamai_ghost = server_2d[:len(akamai_ghost)]
apache = server_2d[len(akamai_ghost):len(akamai_ghost) + len(apache)]
router = server_2d[len(akamai_ghost) + len(apache):]
certificate = trusted + untrusted
certificate_2d = pca(certificate)
trusted = certificate_2d[:len(trusted)]
untrusted = certificate_2d[len(trusted):]
print("ploting....")
plt.figure(figsize = (24, 8))
plt.subplot(1, 3, 1)
plt.title("location")
plt.scatter(china[:, 0], china[:, 1], color='r', label='China')
plt.scatter(germany[:, 0], germany[:, 1], color = 'b', label='Germany')
plt.scatter(japan[:, 0], japan[:, 1], color = 'y', label='Japan')
plt.legend()
plt.subplot(1, 3, 2)
plt.title("servers")
plt.scatter(akamai_ghost[:, 0], akamai_ghost[:, 1], color = 'r', label = 'AkamaiGHost')
plt.scatter(apache[:, 0], apache[:, 1], color = 'b', label = 'Apache')
plt.scatter(router[:, 0], router[:, 1], color = 'y', label = 'Router')
plt.legend()
plt.subplot(1, 3, 3)
plt.title("certificate")
plt.scatter(trusted[:, 0], trusted[:, 1], color = 'r', label = 'Trusted certificate')
plt.scatter(untrusted[:, 0], untrusted[:, 1], color = 'b', label = 'Untrusted certificate')
plt.legend()
plt.savefig(os.path.join(SAVE_PATH, "cluster.png"))
