def __init__(self):
"""
- create a server
- generate training data for classification and regression models based on some function
- train the classifiers and regressors
- store the original function, the ml predictor and training data together for each instance
- load server with a set of trained ml models
- be able to compare an adversarys model to the correct function and initial model
- calculate errors / different error types
"""
self.server = Server()
self.adversary = Adversary()
self.models = {} # {"Name": {"kernel_type": None, "predictor_type": None, "training_data": {"X": None, "y": None}, "original_model": None, "extracted_model": None}}
def __init__(self, network_creator, environment_creator, args):
super(PAACLearner, self).__init__(network_creator, environment_creator,
args)
self.workers = args.emulator_workers
self.network_creator = network_creator # record the network creator in order to create good_network later
self.total_rewards = []
self.adversary = Adversary(args)
# state, reward, episode_over, action
self.variables = [(np.asarray(
[emulator.get_initial_state() for emulator in self.emulators],
dtype=np.uint8)), (np.zeros(self.emulator_counts,
dtype=np.float32)),
(np.asarray([False] * self.emulator_counts,
dtype=np.float32)),
(np.zeros((self.emulator_counts, self.num_actions),
dtype=np.float32))]
self.runners = Runners(EmulatorRunner, self.emulators, self.workers,
self.variables)
self.runners.start()
self.shared_states, self.shared_rewards, self.shared_episode_over, self.shared_actions = self.runners.get_shared_variables(
)
self.summaries_op = tf.summary.merge_all()
self.emulator_steps = [0] * self.emulator_counts
self.total_episode_rewards = self.emulator_counts * [0]
self.actions_sum = np.zeros((self.emulator_counts, self.num_actions))
self.y_batch = np.zeros((self.max_local_steps, self.emulator_counts))
self.adv_batch = np.zeros((self.max_local_steps, self.emulator_counts))
self.rewards = np.zeros((self.max_local_steps, self.emulator_counts))
self.states = np.zeros([self.max_local_steps] +
list(self.shared_states.shape),
dtype=np.uint8)
self.actions = np.zeros(
(self.max_local_steps, self.emulator_counts, self.num_actions))
self.values = np.zeros((self.max_local_steps, self.emulator_counts))
self.episodes_over_masks = np.zeros(
(self.max_local_steps, self.emulator_counts))
from adversary import Adversary
def test_binarySearch(ad):
Seq = [0.0, 0.13, 0.24, 0.31, 0.49, 0.56, 0.61, 0.77, 0.84, 0.98]
testItem = [0.0, 0.04, 0.13, 0.156, 0.24, 0.30, 0.31, 0.371, 0.49, 0.491, 0.56, 0.586, 0.61, 0.666, 0.77, 0.79, 0.84, 0.954, 0.98, 0.999]
for i in testItem:
pos = ad.binarySearch(i, Seq)
print(pos)
if __name__ == '__main__':
ad = Adversary("kmeans", [])
test_binarySearch(ad)
def main():
args = getargs()
print(args)
use_cuda = not args.no_cuda and torch.cuda.is_available()
print("using cuda: ", use_cuda)
torch.manual_seed(args.seed)
device = torch.device("cuda:0" if use_cuda else "cpu")
kwargs = {'num_workers': 1, 'pin_memory': True} if use_cuda else {}
dset = PointCloudDataset(args.local)
train_loader = torch.utils.data.DataLoader(dset,
batch_size=args.batch_size,
shuffle=True,
**kwargs)
test_loader = torch.utils.data.DataLoader(dset,
batch_size=args.test_batch_size,
shuffle=True,
**kwargs)
k = args.k
encoder = Encoder(k).to(device)
decoder = Decoder(k).to(device)
adversary = Adversary(k).to(device)
if not args.save_model_to: Error("Must have a model to save to!")
if args.load_model_from:
encoder.load_state_dict(
torch.load("./models/" + args.load_model_from + "-encoder.pt"))
decoder.load_state_dict(
torch.load("./models/" + args.load_model_from + "-decoder.pt"))
adversary.load_state_dict(
torch.load("./models/" + args.load_model_from + "-adversary.pt"))
else:
params = list(adversary.parameters()) + list(
encoder.parameters()) + list(decoder.parameters())
optimizer = optim.Adam(params, lr=args.lr)
## Visualize one batch of training data
dataiter = iter(train_loader)
runner = RunModel(args)
for epoch in range(args.epochs):
for i, (data, encoding) in enumerate(
runner.train(args, encoder, decoder, adversary, device,
train_loader, optimizer, epoch)):
with torch.no_grad():
# concat = torch.cat((data, encoding), 0)
if i % args.log_interval == 0:
im1 = random.randint(0, args.batch_size - 1)
im2 = random.randint(0, args.batch_size - 1)
imshow(args, data[im1, 0, :, :, :],
data[im2, 0, :, :, :], encoding[im1,
0, :, :, :],
encoding[im2, 0, :, :, :], epoch, i)
imshow(args, data[im1, 0, :, :, :],
data[im2, 0, :, :, :], encoding[im1,
0, :, :, :],
encoding[im2, 0, :, :, :], epoch, i)
imshow(
args, data[im1, 0, :, :, :], encoding[im1,
0, :, :, :],
torch.abs(encoding[im1, 0, :, :, :] -
data[im1, 0, :, :, :]),
torch.abs(encoding[im2, 0, :, :, :] -
data[im2, 0, :, :, :]), epoch,
i + args.batch_size)
save_2d_proj(data[im1, 0, :, :, :], args.save_model_to,
"Train", epoch, i, "original")
save_2d_proj(encoding[im1,
0, :, :, :], args.save_model_to,
"Train", epoch, i, "encoded")
runner.visualize()
runner.test(args, encoder, decoder, adversary, device, test_loader,
args.epochs)
if args.save_model_to:
print("saving model")
torch.save(encoder.state_dict(),
"./models/" + args.save_model_to + "-encoder.pt")
torch.save(decoder.state_dict(),
"./models/" + args.save_model_to + "-decoder.pt")
torch.save(adversary.state_dict(),
"./models/" + args.save_model_to + "-adversary.pt")
print("quantify")
if args.quantify:
quantify_results(encoder, decoder, adversary, test_loader, args,
device)
print(result)
# 4
detect_single_charactor_xor()
# # 5
out = repeating_xor(
b"Burning 'em, if you ain't quick and nimble\nI go crazy when I hear a cymbal",
b"ICE")
assert (a2b_hex(
'0b3637272a2b2e63622c2e69692a23693a2a3c6324202d623d63343c2a26226324272765272a282b2f20430a652e2c652a3124333a653e2b2027630c692b20283165286326302e27282f'
) == out)
#
# 5
out = hamming_distance(b'this is a test', b'wokka wokka!!!')
assert (out == 37)
adversary = Adversary()
with open('6.txt') as f:
contents = base64.b64decode(f.read())
print(adversary.crack_repeating_xor(contents))
print(decrypt_aes_ecb().decode('utf-8'))
with open("8.txt") as f:
contents = []
for line in f:
contents.append(line.strip())
print(detect_aes_ecb(contents))
print("Success")
def main():
args = init_args()
if os.path.exists(args.output_dir) and os.listdir(
args.output_dir
) and args.do_train and not args.overwrite_output_dir:
raise ValueError(
"Output directory ({}) already exists and is not empty. Use --overwrite_output_dir to overcome."
.format(args.output_dir))
# Setup CUDA, GPU & distributed training
if args.local_rank == -1 or args.no_cuda:
device = torch.device("cuda" if torch.cuda.is_available()
and not args.no_cuda else "cpu")
args.n_gpu = torch.cuda.device_count()
else: # Initializes the distributed backend which will take care of sychronizing nodes/GPUs
torch.cuda.set_device(args.local_rank)
device = torch.device("cuda", args.local_rank)
os.environ['MASTER_ADDR'] = args.MASTER_ADDR
os.environ['MASTER_PORT'] = args.MASTER_PORT
torch.distributed.init_process_group(backend='nccl',
rank=args.local_rank,
world_size=1)
args.n_gpu = 1
args.device = device
# Setup logging
logging.basicConfig(
format='%(asctime)s - %(levelname)s - %(name)s - %(message)s',
datefmt='%m/%d/%Y %H:%M:%S',
level=logging.INFO if args.local_rank in [-1, 0] else logging.WARN)
# not using 16-bits training
logger.warning(
"Process rank: %s, device: %s, n_gpu: %s, distributed training: %s, 16-bits training: False",
args.local_rank, device, args.n_gpu, bool(args.local_rank != -1))
# Set seed
set_seed(args)
# Prepare task
args.task_name = args.task_name.lower()
if args.task_name not in processors:
raise ValueError("Task not found: %s" % args.task_name)
processor = processors[args.task_name]()
args.output_mode = output_modes[args.task_name]
label_list = processor.get_labels(args.tagging_schema)
num_labels = len(label_list)
normal_labels = processor.get_normal_labels(args.tagging_schema)
num_normal_labels = len(normal_labels)
sent_labels = ABSAProcessor.get_sentiment_labels()
num_sent_labels = len(sent_labels)
# initialize the pre-trained model
args.model_type = args.model_type.lower()
config_class, model_class, tokenizer_class = MODEL_CLASSES[args.model_type]
config = config_class.from_pretrained(
args.config_name if args.config_name else args.model_name_or_path,
num_labels=num_labels,
finetuning_task=args.task_name)
tokenizer = AutoTokenizer.from_pretrained(args.model_name_or_path,
cache_dir='./cache')
config.absa_type = args.absa_type
config.tfm_mode = args.tfm_mode
config.fix_tfm = args.fix_tfm
config.num_normal_labels = num_normal_labels
config.num_sent_labels = num_sent_labels
config.ts_vocab = {label: i for i, label in enumerate(label_list)}
config.ote_vocab = {label: i for i, label in enumerate(normal_labels)}
config.sent_vocab = {label: i for i, label in enumerate(sent_labels)}
config.device = "cuda" if torch.cuda.is_available(
) and not args.no_cuda else "cpu"
config.output_hidden_states = True
config.model_name_or_path = args.model_name_or_path
if args.gen_adv_from_path:
# Generate adversarial examples
modes = ['train', 'dev', 'test']
for mode in modes:
model = model_class.from_pretrained(args.gen_adv_from_path).to(
args.device)
train_dataset, train_evaluate_label_ids, examples, imp_words = load_and_cache_examples(
args, args.task_name, tokenizer, mode=mode, model=model)
adversary = Adversary(args, model)
adv_examples = []
sz = 64
for _ in trange(len(examples) // sz + 1):
if len(examples) == 0:
continue
adv_examples.extend(
adversary.generate_adv_examples(examples[:sz],
imp_words[:sz], tokenizer))
examples = examples[sz:]
imp_words = imp_words[sz:]
adv_dataset = convert_to_dataset(args, adv_examples, tokenizer)
output_dir = f'{args.task_name}_adv'
if not os.path.exists(output_dir):
os.makedirs(output_dir)
torch.save(adv_dataset, f'{output_dir}/{mode}.pth')
torch.save(adv_examples, f'{output_dir}/{mode}-examples.pth')
exit(0)
if args.load_model:
print('Loading model from:', args.load_model)
model = model_class.from_pretrained(args.load_model, config=config)
else:
model = model_class.from_pretrained(args.model_name_or_path,
config=config,
cache_dir='./cache')
print('Loading model from:', args.model_name_or_path)
# Distributed and parallel training
model.to(args.device)
if args.local_rank != -1:
model = torch.nn.parallel.DistributedDataParallel(
model,
device_ids=[args.local_rank],
output_device=args.local_rank,
find_unused_parameters=True)
elif args.n_gpu > 1:
model = torch.nn.DataParallel(model)
# Training
if args.do_train:
# Create output directory if needed
if not os.path.exists(args.output_dir) and args.local_rank in [-1, 0]:
os.mkdir(args.output_dir)
# Store model configuration with results
shutil.copyfile('absa_layer.py', args.output_dir + '/absa_layer.py')
# Store training configuration
shutil.copyfile('train.sh', args.output_dir + '/train.sh')
if args.do_adv:
# Store adv training config
shutil.copyfile('main.py', args.output_dir + '/main.py')
train_dataset, train_evaluate_label_ids, examples, imp_words = load_and_cache_examples(
args, args.task_name, tokenizer, mode='train', model=model)
global_step, tr_loss = train(args, train_dataset, model, tokenizer)
model_to_save = model.module if hasattr(model, 'module') else model
model_to_save.save_pretrained(args.output_dir)
tokenizer.save_pretrained(args.output_dir)
# Good practice: save your training arguments together with the trained model
# save the model configuration
torch.save(args, os.path.join(args.output_dir, 'training_args.bin'))
# Load a trained model and vocabulary that you have fine-tuned
model = model_class.from_pretrained(args.output_dir)
tokenizer = tokenizer_class.from_pretrained(args.output_dir)
model.to(args.device)
# Validation
results = {}
best_f1 = -999999.0
best_checkpoint = None
checkpoints = []
if args.eval_all_checkpoints:
checkpoints = os.listdir(args.output_dir)
checkpoints.sort()
logger.info("Perform validation on the following checkpoints: %s",
checkpoints)
test_results = {}
steps = []
for checkpoint in checkpoints:
global_step = checkpoint.split('-')[-1] if len(checkpoints) > 1 else ""
if checkpoint.split('-')[0] != 'checkpoint':
continue
if args.pred_checkpoint and args.pred_checkpoint != global_step:
continue
steps.append(global_step)
set_seed(args)
model = model_class.from_pretrained(f'{args.output_dir}/{checkpoint}')
model.to(args.device)
dev_result = evaluate(args,
model,
tokenizer,
mode='dev',
prefix=global_step)
# regard the micro-f1 as the criteria of model selection
if int(global_step) > 1000 and dev_result['micro-f1'] > best_f1:
best_f1 = dev_result['micro-f1']
best_checkpoint = checkpoint
dev_result = dict(
(k + '_{}'.format(global_step), v) for k, v in dev_result.items())
results.update(dev_result)
test_result = evaluate(args,
model,
tokenizer,
mode='test',
prefix=global_step)
test_result = dict(
(k + '_{}'.format(global_step), v) for k, v in test_result.items())
test_results.update(test_result)
best_ckpt_string = "\nThe best checkpoint is %s" % best_checkpoint
logger.info(best_ckpt_string)
dev_f1_values, dev_loss_values = [], []
for k in results:
v = results[k]
if 'micro-f1' in k:
dev_f1_values.append((k, v))
if 'eval_loss' in k:
dev_loss_values.append((k, v))
test_f1_values, test_loss_values = [], []
for k in test_results:
v = test_results[k]
if 'micro-f1' in k:
test_f1_values.append((k, v))
if 'eval_loss' in k:
test_loss_values.append((k, v))
log_file_path = '%s/log.txt' % args.output_dir
log_file = open(log_file_path, 'a')
log_file.write("\tValidation:\n")
for (test_f1_k,
test_f1_v), (test_loss_k,
test_loss_v), (dev_f1_k,
dev_f1_v), (dev_loss_k,
dev_loss_v) in zip(
test_f1_values,
test_loss_values,
dev_f1_values,
dev_loss_values):
global_step = int(test_f1_k.split('_')[-1])
if not args.overfit and global_step <= 1000:
continue
print('test-%s: %.5lf, test-%s: %.5lf, dev-%s: %.5lf, dev-%s: %.5lf' %
(test_f1_k, test_f1_v, test_loss_k, test_loss_v, dev_f1_k,
dev_f1_v, dev_loss_k, dev_loss_v))
validation_string = '\t\tdev-%s: %.5lf, dev-%s: %.5lf' % (
dev_f1_k, dev_f1_v, dev_loss_k, dev_loss_v)
log_file.write(validation_string + '\n')
n_times = args.max_steps // args.save_steps + 1
for step in steps:
log_file.write('\tStep %s:\n' % step)
precision = test_results['precision_%s' % step]
recall = test_results['recall_%s' % step]
micro_f1 = test_results['micro-f1_%s' % step]
macro_f1 = test_results['macro-f1_%s' % step]
log_file.write(
'\t\tprecision: %.4lf, recall: %.4lf, micro-f1: %.4lf, macro-f1: %.4lf\n'
% (precision, recall, micro_f1, macro_f1))
log_file.write("\tBest checkpoint: %s\n" % best_checkpoint)
log_file.write('******************************************\n')
log_file.close()
raw_filenames = ["raw_training_text", "raw_valid_text", "raw_test_text"]
suffix = ".txt"
# In[ ]:
dataset_dict = load_pickle(join(path, dataset_dict_filename))
index2token = {index: token for (token, index, _, _) in dataset_dict}
token2index = {token: index for (token, index, _, _) in dataset_dict}
vocab = list(token2index.keys())
indices = list(token2index.values())
vocab_file = join(path, "vocab.txt")
if not exists(vocab_file):
write_lines(vocab_file, vocab)
# In[ ]:
adversary = Adversary(vocab)
# In[ ]:
for filename in raw_filenames:
file = join(path, "raw", filename + suffix)
dialogues = read_lines(file)
adv_dialogues = []
for dialogue in tqdm(dialogues):
adv_dialogue = adversary.apply_strategy(strategy, dialogue)
adv_dialogues.append(adv_dialogue)
target_file = join(path, "raw", filename + "." + strategy + suffix)
write_lines(target_file, adv_dialogues)
class Supervisor:
"""
This Class delegates the whole experiment.
It keeps track of actual functions that are mimicked by ML, of the ML models, and the stolen Models
in order to provide overview and comparison between those stages.
This class is used to calculate errors and optimize the extraction algorithm.
It can be seen as an omniscient class that knows of everything going on in this program
"""
def __init__(self):
"""
- create a server
- generate training data for classification and regression models based on some function
- train the classifiers and regressors
- store the original function, the ml predictor and training data together for each instance
- load server with a set of trained ml models
- be able to compare an adversarys model to the correct function and initial model
- calculate errors / different error types
"""
self.server = Server()
self.adversary = Adversary()
self.models = {} # {"Name": {"kernel_type": None, "predictor_type": None, "training_data": {"X": None, "y": None}, "original_model": None, "extracted_model": None}}
def add_model(self, name, X, y, predictor_type, kernel_type):
"""
Trains an ML model on the given data, adds it to the class internal database, and adds it to the server
:param name: Name of the model
:param X: TrainingData X
:param y: TrainingData Targets
:param predictor_type: SVM or SVR
:param kernel_type: kernel to use.
:return:
"""
if kernel_type == "quadratic":
kernel_type = "poly"
deg = 2
else:
deg = 3
if predictor_type == "SVM":
original_model = svm.SVC(kernel=kernel_type, degree=deg)
original_model.fit(X, y)
elif predictor_type == "SVR":
original_model = svm.SVR(kernel=kernel_type, degree=deg)
original_model.fit(X, y)
else:
raise ValueError
self.models[name] = {"training_data": {"X": X, "y": y},
"kernel_type": kernel_type,
"predictor_type": predictor_type,
"original_model": original_model,
"extracted_model": None}
self.server.add_predictor(predictor_type, name, original_model)
return self.models[name]
def get_models(self):
return self.models
def compare_predictions(self, name, data, correct_results=None, verbose=False):
"""
Compares predictions on dataset "data" between original and extracted model. Returns error percentage for svm
and mean, mean squared and relative mean squared error for svr.
:param name:
:param data:
:param correct_results:
:param verbose:
:return:
"""
prediction_errors = []
original_model = self.models[name]["original_model"]
extracted_model = self.models[name]["extracted_model"]
prediction_type = self.models[name]["predictor_type"]
error_count = 0
original_predictions = []
for index, datum in enumerate(data):
original_prediction = original_model.predict([datum])
extracted_prediction = extracted_model.predict([datum])
original_predictions.append(original_prediction[0])
error = abs(original_prediction[0]-extracted_prediction[0])
prediction_errors.append(error)
if error != 0 and prediction_type == "SVM":
error_count += 1
if verbose:
if error == 1:
pass
#print(original_prediction)
#print(extracted_prediction)
#print(datum.T[0],",",datum.T[1])
elif 1==2:
print("data:", datum)
if correct_results is not None:
print("correct: ", correct_results[index])
print("original prediction: ", original_prediction[0])
print("extracted prediction:", extracted_prediction[0])
print("error ", error)
print("------------------------------")
if prediction_type == "SVM":
error_percentage = error_count*100/len(data)
print("Total false Predictions", error_count, "out of", len(data))
print("False Prediction Percentage:", error_percentage, "%")
return error_percentage, 0, 0
elif prediction_type == "SVR":
mean_original = sum(original_predictions) / len(original_predictions)
mean_error = sum(prediction_errors) / len(prediction_errors)
mse = sum(map(lambda x: x ** 2, prediction_errors)) / len(prediction_errors)
rmse = sum(map(lambda x: x ** 2, prediction_errors)) / sum(
map(lambda x: (mean_original - x) ** 2, original_predictions))
print("----------------------------------------")
print("Mean Error:", mean_error)
print("MSE: ", mse)
print("RMSE: ", rmse)
return mean_error, mse, rmse
def attack_model(self, name, kernel_type, attack_type, dimension, query_budget, dataset=None, roundsize=5, test_set=None):
if attack_type == "lowd-meek" or attack_type == "extraction":
print("[*] Original Model Parameters")
if kernel_type == "linear":
print("w ", self.models[name]["original_model"].coef_)
print("a ", -self.models[name]["original_model"].coef_[0][0]/self.models[name]["original_model"].coef_[0][1])
print("b ", -self.models[name]["original_model"].intercept_[0] / self.models[name]["original_model"].coef_[0][1])
else:
print("a ", self.models[name]["original_model"].dual_coef_)
#print("sv", self.models[name]["original_model"].support_vectors_)
predictor_type = self.models[name]["predictor_type"]
if kernel_type is None and attack_type == "agnostic" and test_set is not None:
self.models[name]["extracted_model"] = self.adversary.kernel_agnostic_attack(self.server, name,
predictor_type, dimension,
query_budget, dataset, test_set)
else:
self.models[name]["extracted_model"] = self.adversary.attack(self.server, name, predictor_type,
kernel_type, attack_type, dimension,
query_budget, dataset=dataset, roundsize=roundsize)
return self.models[name]
def attack_with_metrics(self, name, kernel_type, attack_type, dimension, query_budget, dataset=None, roundsize=5, test_set=None):
start_time = time.time()
self.attack_model(name, kernel_type, attack_type, dimension, query_budget,
dataset=dataset, roundsize=roundsize, test_set=test_set)
queries = self.adversary.get_last_query_count()
run_time = time.time() - start_time
return run_time, queries, self.models[name]["extracted_model"]
def suggest_attack_type(self, predictor_type, kernel_type):
if predictor_type == "SVM" or predictor_type == "svm" or predictor_type == "SVC" or predictor_type == "svc":
if kernel_type == "linear":
return "lowd-meek"
elif kernel_type == "polynomial":
return "lowd-meek"
else:
return "adaptive retraining"
elif predictor_type == "SVR" or predictor_type == "svr":
if kernel_type == "linear":
return "extraction"
elif kernel_type == "polynomial":
return "extraction"
else:
return "adaptive retraining"
else:
raise ValueError
def get_error_on_queries(self, error_type, name, kernel_type, attack_type, dimension, budget_factor_alpha_list, training_data, test_data, roundsize=5):
if not isinstance(test_data, list):
test_data = test_data.tolist()
if not isinstance(training_data, list):
training_data = training_data.tolist()
x_axis = budget_factor_alpha_list
y_axis = []
for budget_factor_alpha in budget_factor_alpha_list:
print("Alpha:", budget_factor_alpha)
query_budget = math.ceil(budget_factor_alpha * (dimension + 1))
run_time, queries, model = self.attack_with_metrics(name, kernel_type, attack_type, dimension, query_budget, training_data, roundsize=roundsize)
print("Queries", queries, " Runtime", run_time)
errors = self.compare_predictions(name, test_data, verbose=False)
if self.models[name]["predictor_type"] == "SVR":
if error_type == "mean":
error_index = 0
elif error_type == "mse":
error_index = 1
elif error_type == "rmse":
error_index = 2
y_axis.append(errors[error_index])
else:
y_axis.append(errors[0])
return x_axis, y_axis
def plot_error_on_queries(self, error_type, name, kernel_type, attack_type, dimension, budget_factor_alpha_list, training_data, test_data, roundsize=5):
x_axis, y_axis = self.get_error_on_queries(error_type, name, kernel_type, attack_type, dimension, budget_factor_alpha_list, training_data, test_data, roundsize=roundsize)
[print(i) for i in zip(x_axis, y_axis)]
plt.plot(x_axis, y_axis)
plt.show()
return plt, x_axis, y_axis
def main(args):
robot = RobotEnv()
U.make_session(num_cpu=1).__enter__()
set_global_seeds(args.seed)
def policy_fn(name, ob_space, ac_space, reuse=False):
return mlp_policy.MlpPolicy(name=name, ob_space=ob_space, ac_space=ac_space,
reuse=reuse, hid_size=args.policy_hidden_size, num_hid_layers=3)
task_name = U_.get_task_name(args)
args.checkpoint_dir = osp.join(args.checkpoint_dir, task_name)
args.log_dir = osp.join(args.log_dir, task_name)
dataset = UR_Dset(expert_path=args.expert_path)
ob_shape = dataset.get_ob_shape()
ac_shape = dataset.get_ac_shape()
robot.set_ob_shape(ob_shape)
robot.set_ac_shape(ac_shape)
'''
pi = bc_learn(args.evaluate,
robot,
policy_fn,
dataset,
ckpt_dir=args.checkpoint_dir,
log_dir=args.log_dir,
task_name=task_name,
verbose=True)
from time import sleep
ob = robot.reset()
sleep(2)
# testing directly mapping
for _ in tqdm(range(1000)):
_, acs = dataset.get_next_batch(1, 'train')
_ = robot.step(acs[0])
# sleep(0.5)
# testing policy
for _ in tqdm(range(1000)):
ac, _ = pi.act(True, ob)
ob = robot.step(ac)
# sleep(0.5)
'''
if args.task == 'train':
reward_giver = Adversary(robot, args.adversary_hidden_size)
trpo_train(robot,
args.seed,
policy_fn,
reward_giver,
dataset,
args.g_step,
args.d_step,
args.policy_entcoeff,
args.num_timesteps,
args.save_per_iter,
args.checkpoint_dir,
args.log_dir,
task_name
)
elif args.task == 'evaluate':
runner(robot,
policy_fn,
args.load_model_path,
timesteps_per_batch=1024,
number_trajs=10,
stochastic_policy=args.stochastic_policy,
save=args.save_sample
)
else:
raise NotImplementedError
# print (list(zip(Z,list(X))))
a = list(zip(Z, X))
# print (a[0][0])
# print (a[0][1])
sZ, sX = zip(*sorted(a, key=lambda x: x[0]))
# apply PAV
PAVY = pav(sY)
for i in range(len(sX)):
w = w + (sY[i] - PAVY[i]) * sX[i] / len(sX)
return [w, sZ, PAVY]
T = 200
D = 5
R = 5
adv = Adversary(D, T, R)
E = 0.0
w = np.array([0 for x in range(D)])
for t in range(1, T):
X, Z, Y = adv.train_data[0][:t], adv.train_data[1][:t], adv.train_data[
2][:t]
w, sZ, PAVY = slisotron(w, X, Y)
# get the error of t'th point
x, z, y = adv.train_data[0][t], adv.train_data[1][t], adv.train_data[2][t]
# search z in sZ and get the adjacent point
pos = bisect.bisect_left(sZ, z)
if pos == 0:
err = math.pow(PAVY[0] - y, 2)
E = E + err
print(t, err, E)
elif pos == len(sZ):
def __init__(self, config, features, labels, args, evaluate=False):
# Build the computational graph
arch = Network.sequence_deep_conv
if args.architecture == 'deep_conv':
arch = Network.sequence_deep_conv
elif args.architecture == 'recurrent':
arch = Network.birnn_dynamic
elif args.architecture == 'simple_conv':
arch = Network.sequence_conv2d
elif args.architecture == 'conv_projection':
arch = Network.conv_projection
self.global_step = tf.Variable(0, trainable=False)
self.handle = tf.placeholder(tf.string, shape=[])
self.training_phase = tf.placeholder(tf.bool)
self.rnn_keep_prob = tf.placeholder(tf.float32)
self.features_placeholder = tf.placeholder(tf.float32, [features.shape[0], features.shape[1]])
self.labels_placeholder = tf.placeholder(tf.int32, labels.shape)
self.test_features_placeholder = tf.placeholder(tf.float32)
self.test_labels_placeholder = tf.placeholder(tf.int32)
steps_per_epoch = int(self.features_placeholder.get_shape()[0])//config.batch_size
if config.use_adversary and not evaluate:
self.pivots_placeholder = tf.placeholder(tf.float32)
self.pivot_labels_placeholder = tf.placeholder(tf.int32)
self.test_pivots_placeholder = tf.placeholder(tf.float32)
self.test_pivot_labels_placeholder = tf.placeholder(tf.int32)
train_dataset = Data.load_dataset_adversary(self.features_placeholder, self.labels_placeholder,
self.pivots_placeholder, self.pivot_labels_placeholder, config.batch_size)
test_dataset = Data.load_dataset_adversary(self.test_features_placeholder, self.test_labels_placeholder,
self.test_pivots_placeholder, self.test_pivot_labels_placeholder, config_test.batch_size, test=True)
else:
train_dataset = Data.load_dataset(self.features_placeholder, self.labels_placeholder,
config.batch_size)
test_dataset = Data.load_dataset(self.test_features_placeholder, self.test_labels_placeholder,
config_test.batch_size, test=True)
val_dataset = Data.load_dataset(self.features_placeholder, self.labels_placeholder, config.batch_size, evaluate=True)
self.iterator = tf.contrib.data.Iterator.from_string_handle(self.handle,
train_dataset.output_types,
train_dataset.output_shapes)
self.train_iterator = train_dataset.make_initializable_iterator()
self.test_iterator = test_dataset.make_initializable_iterator()
self.val_iterator = val_dataset.make_initializable_iterator()
# embedding_encoder = tf.get_variable('embeddings', [config.features_per_particle, config.embedding_dim])
if config.use_adversary:
self.example, self.labels, self.pivots, self.pivot_labels = self.iterator.get_next()
if len(config.pivots) == 1:
self.pivots = tf.expand_dims(self.pivots, axis=1)
self.pivot_labels = tf.expand_dims(self.pivot_labels, axis=1)
print(self.pivots.get_shape())
else:
self.example, self.labels = self.iterator.get_next()
if evaluate:
# embeddings = tf.nn.embedding_lookup(embedding_encoder, ids=self.example)
with tf.variable_scope('classifier') as scope:
self.logits = arch(self.example, config, self.training_phase)
self.softmax, self.pred = tf.nn.softmax(self.logits)[:,1], tf.argmax(self.logits, 1)
self.ema = tf.train.ExponentialMovingAverage(decay=config.ema_decay, num_updates=self.global_step)
return
# embeddings = tf.nn.embedding_lookup(embedding_encoder, ids=self.example)
with tf.variable_scope('classifier') as scope:
self.logits = arch(self.example, config, self.training_phase)
self.softmax, self.pred = tf.nn.softmax(self.logits), tf.argmax(self.logits, 1)
epoch_bounds = [64, 128, 256, 420, 512, 720, 1024]
lr_values = [1e-3, 4e-4, 1e-4, 6e-5, 1e-5, 6e-6, 1e-6, 2e-7]
learning_rate = tf.train.piecewise_constant(self.global_step, boundaries=[s*steps_per_epoch for s in
epoch_bounds], values=lr_values)
if config.use_adversary:
adv = Adversary(config,
classifier_logits=self.logits,
labels=self.labels,
pivots=self.pivots,
pivot_labels=self.pivot_labels,
training_phase=self.training_phase,
predictor_learning_rate=learning_rate,
args=args)
self.cost = adv.predictor_loss
self.adv_loss = adv.adversary_combined_loss
self.total_loss = adv.total_loss
self.predictor_train_op = adv.predictor_train_op
self.adversary_train_op = adv.adversary_train_op
self.joint_step, self.ema = adv.joint_step, adv.ema
self.joint_train_op = adv.joint_train_op
else:
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
self.cross_entropy = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=self.logits,
labels=self.labels)
self.cost = tf.reduce_mean(self.cross_entropy)
with tf.control_dependencies(update_ops):
# Ensures that we execute the update_ops before performing the train_step
if args.optimizer=='adam':
self.opt_op = tf.train.AdamOptimizer(learning_rate).minimize(self.cost, global_step=self.global_step)
elif args.optimizer=='momentum':
self.opt_op = tf.train.MomentumOptimizer(learning_rate, config.momentum,
use_nesterov=True).minimize(self.cost, global_step=self.global_step)
self.ema = tf.train.ExponentialMovingAverage(decay=config.ema_decay, num_updates=self.global_step)
maintain_averages_op = self.ema.apply(tf.trainable_variables())
with tf.control_dependencies(update_ops+[self.opt_op]):
self.train_op = tf.group(maintain_averages_op)
self.str_accuracy, self.update_accuracy = tf.metrics.accuracy(self.labels, self.pred)
correct_prediction = tf.equal(self.labels, tf.cast(self.pred, tf.int32))
_, self.auc_op = tf.metrics.auc(predictions=self.pred, labels=self.labels, num_thresholds=1024)
self.accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))
tf.summary.scalar('accuracy', self.accuracy)
tf.summary.scalar('learning_rate', learning_rate)
tf.summary.scalar('cost', self.cost)
tf.summary.scalar('auc', self.auc_op)
self.merge_op = tf.summary.merge_all()
self.train_writer = tf.summary.FileWriter(
os.path.join(directories.tensorboard, '{}_train_{}'.format(args.name, time.strftime('%d-%m_%I:%M'))), graph=tf.get_default_graph())
self.test_writer = tf.summary.FileWriter(
os.path.join(directories.tensorboard, '{}_test_{}'.format(args.name, time.strftime('%d-%m_%I:%M'))))
class PAACLearner(ActorLearner):
def __init__(self, network_creator, environment_creator, args):
super(PAACLearner, self).__init__(network_creator, environment_creator,
args)
self.workers = args.emulator_workers
self.network_creator = network_creator # record the network creator in order to create good_network later
self.total_rewards = []
self.adversary = Adversary(args)
# state, reward, episode_over, action
self.variables = [(np.asarray(
[emulator.get_initial_state() for emulator in self.emulators],
dtype=np.uint8)), (np.zeros(self.emulator_counts,
dtype=np.float32)),
(np.asarray([False] * self.emulator_counts,
dtype=np.float32)),
(np.zeros((self.emulator_counts, self.num_actions),
dtype=np.float32))]
self.runners = Runners(EmulatorRunner, self.emulators, self.workers,
self.variables)
self.runners.start()
self.shared_states, self.shared_rewards, self.shared_episode_over, self.shared_actions = self.runners.get_shared_variables(
)
self.summaries_op = tf.summary.merge_all()
self.emulator_steps = [0] * self.emulator_counts
self.total_episode_rewards = self.emulator_counts * [0]
self.actions_sum = np.zeros((self.emulator_counts, self.num_actions))
self.y_batch = np.zeros((self.max_local_steps, self.emulator_counts))
self.adv_batch = np.zeros((self.max_local_steps, self.emulator_counts))
self.rewards = np.zeros((self.max_local_steps, self.emulator_counts))
self.states = np.zeros([self.max_local_steps] +
list(self.shared_states.shape),
dtype=np.uint8)
self.actions = np.zeros(
(self.max_local_steps, self.emulator_counts, self.num_actions))
self.values = np.zeros((self.max_local_steps, self.emulator_counts))
self.episodes_over_masks = np.zeros(
(self.max_local_steps, self.emulator_counts))
@staticmethod
def choose_next_actions(network, num_actions, states, session):
network_output_v, network_output_pi = session.run(
[network.output_layer_v, network.output_layer_pi],
feed_dict={network.input_ph: states})
action_indices = PAACLearner.__sample_policy_action(network_output_pi)
new_actions = np.eye(num_actions)[action_indices]
return new_actions, network_output_v, network_output_pi
def __choose_next_actions(self, states):
return PAACLearner.choose_next_actions(self.network, self.num_actions,
states, self.session)
@staticmethod
def __sample_policy_action(probs):
"""
Sample an action from an action probability distribution output by
the policy network.
"""
# Subtract a tiny value from probabilities in order to avoid
# "ValueError: sum(pvals[:-1]) > 1.0" in numpy.multinomial
probs = probs - np.finfo(np.float32).epsneg
action_indices = [
int(np.nonzero(np.random.multinomial(1, p))[0]) for p in probs
]
return action_indices
def _get_shared(self, array, dtype=c_float):
"""
Returns a RawArray backed numpy array that can be shared between processes.
:param array: the array to be shared
:param dtype: the RawArray dtype to use
:return: the RawArray backed numpy array
"""
shape = array.shape
shared = RawArray(dtype, array.reshape(-1))
return np.frombuffer(shared, dtype).reshape(shape)
def run_policy(self, t):
state_id = self.global_step
self.poisoned_emulators = []
#print('state_id', state_id, 't', t)
self.shared_states = self.adversary.manipulate_states(
state_id, t, self.shared_states)
self.next_actions, readouts_v_t, readouts_pi_t = self.__choose_next_actions(
self.shared_states)
self.next_actions = self.adversary.manipulate_actions(
self.next_actions)
self.actions_sum += self.next_actions
for z in range(self.next_actions.shape[0]):
self.shared_actions[z] = self.next_actions[z]
self.actions[t] = self.next_actions
self.values[t] = readouts_v_t
self.states[t] = self.shared_states
# Start updating all environments with next_actions
self.runners.update_environments()
self.runners.wait_updated()
# Done updating all environments, have new states, rewards and is_over
self.episodes_over_masks[t] = 1.0 - self.shared_episode_over.astype(
np.float32)
def store_rewards(self, t, emulator, actual_reward, episode_over):
actual_reward = self.adversary.poison_reward(emulator, actual_reward,
self.next_actions)
self.total_episode_rewards[emulator] += actual_reward
actual_reward = self.rescale_reward(actual_reward)
self.rewards[t, emulator] = actual_reward
self.emulator_steps[emulator] += 1
if episode_over:
self.total_rewards.append(self.total_episode_rewards[emulator])
episode_summary = tf.Summary(value=[
tf.Summary.Value(
tag='rl/reward',
simple_value=self.total_episode_rewards[emulator]),
tf.Summary.Value(tag='rl/episode_length',
simple_value=self.emulator_steps[emulator]),
])
self.summary_writer.add_summary(episode_summary, self.global_step)
self.summary_writer.flush()
self.total_episode_rewards[emulator] = 0
self.emulator_steps[emulator] = 0
self.actions_sum[emulator] = np.zeros(self.num_actions)
def calculate_estimated_return(self):
nest_state_value = self.session.run(
self.network.output_layer_v,
feed_dict={self.network.input_ph: self.shared_states})
estimated_return = np.copy(nest_state_value)
for t in reversed(range(self.max_local_steps)):
estimated_return = self.rewards[
t] + self.gamma * estimated_return * self.episodes_over_masks[t]
self.y_batch[t] = np.copy(estimated_return)
self.adv_batch[t] = estimated_return - self.values[t]
def update_networks(self):
flat_states = self.states.reshape(
[self.max_local_steps * self.emulator_counts] +
list(self.shared_states.shape)[1:])
flat_y_batch = self.y_batch.reshape(-1)
flat_adv_batch = self.adv_batch.reshape(-1)
flat_actions = self.actions.reshape(
self.max_local_steps * self.emulator_counts, self.num_actions)
lr = self.get_lr()
feed_dict = {
self.network.input_ph: flat_states,
self.network.critic_target_ph: flat_y_batch,
self.network.selected_action_ph: flat_actions,
self.network.adv_actor_ph: flat_adv_batch,
self.learning_rate: lr
}
_, summaries = self.session.run([self.train_step, self.summaries_op],
feed_dict=feed_dict)
self.summary_writer.add_summary(summaries, self.global_step)
self.summary_writer.flush()
def train(self):
"""
Main actor learner loop for parallel advantage actor critic learning.
"""
self.global_step = self.init_network()
self.last_saving_step = self.global_step
logging.debug("Starting training at Step {}".format(self.global_step))
counter = 0
global_start = self.global_step
start_time = time.time()
print("global_step: ", self.global_step)
while self.global_step < self.max_global_steps:
loop_start_time = time.time()
for t in range(self.max_local_steps):
self.run_policy(t)
for e, (actual_reward, episode_over) in enumerate(
zip(self.shared_rewards, self.shared_episode_over)):
self.global_step += 1
self.store_rewards(t, e, actual_reward, episode_over)
self.calculate_estimated_return()
self.update_networks()
counter += 1
if counter % (2048 / self.emulator_counts) == 0:
curr_time = time.time()
global_steps = self.global_step
last_ten = 0.0 if len(self.total_rewards) < 1 else np.mean(
self.total_rewards[-10:])
logging.info(
"Ran {} steps, at {} steps/s ({} steps/s avg), last 10 rewards avg {}"
.format(
global_steps, self.max_local_steps *
self.emulator_counts / (curr_time - loop_start_time),
(global_steps - global_start) /
(curr_time - start_time), last_ten))
print(datetime.datetime.now().strftime("%Y-%b-%d %H:%M"))
print("total_poison: ", self.adversary.total_poison)
self.save_vars()
self.cleanup()
with open(os.path.join(self.debugging_folder, 'no_of_poisoned_states'),
'w') as f:
f.write('total_poison: ' + str(self.adversary.total_poison) + '\n')
with open(
os.path.join(self.debugging_folder, 'no_of_poisoned_actions'),
'w') as f:
f.write('target_action: ' +
str(self.adversary.total_target_actions) + '\n')
f.write('poison_distribution: ' +
str(self.adversary.poison_distribution) + '\n')
if self.adversary.attack_method == 'untargeted':
with open(
os.path.join(self.debugging_folder,
'no_of_poisoned_rewards_to_one'), 'w') as f:
f.write('total times we give reward 1: ' +
str(self.adversary.total_positive_rewards) + '\n')
f.write('total times we give reward -1: ' +
str(self.adversary.total_negative_rewards) + '\n')
else:
with open(
os.path.join(self.debugging_folder,
'no_of_poisoned_rewards_to_one'), 'w') as f:
f.write('total times we give reward 1: ' +
str(self.adversary.total_positive_rewards) + '\n')
def cleanup(self):
super(PAACLearner, self).cleanup()
self.runners.stop()
def __init__(self, config, features, labels, args=None, evaluate=False):
if args is not None: # merge args and config
config_d = dict((n, getattr(config, n)) for n in dir(config)
if not n.startswith('__'))
config_d.update(vars(args))
config = Utils.Struct(**config_d)
# Build the computational graph
arch = functools.partial(Network.dense_network, actv=tf.nn.elu)
sequential = False
self.global_step = tf.Variable(0, trainable=False)
self.MINE_step = tf.Variable(0, trainable=False)
self.handle = tf.placeholder(tf.string, shape=[])
self.training_phase = tf.placeholder(tf.bool)
self.rnn_keep_prob = tf.placeholder(tf.float32)
self.features_placeholder = tf.placeholder(
tf.float32, [features.shape[0], features.shape[1]])
self.labels_placeholder = tf.placeholder(tf.int32, labels.shape)
self.test_features_placeholder = tf.placeholder(tf.float32)
self.test_labels_placeholder = tf.placeholder(tf.int32)
self.pivots_placeholder = tf.placeholder(tf.float32)
self.test_pivots_placeholder = tf.placeholder(tf.float32)
if config.use_adversary and not evaluate:
self.pivot_labels_placeholder = tf.placeholder(tf.int32)
self.test_pivot_labels_placeholder = tf.placeholder(tf.int32)
train_dataset = Data.load_dataset(
self.features_placeholder,
self.labels_placeholder,
self.pivots_placeholder,
batch_size=config.batch_size,
sequential=sequential,
adversary=True,
pivot_labels_placeholder=self.pivot_labels_placeholder)
test_dataset = Data.load_dataset(
self.test_features_placeholder,
self.test_labels_placeholder,
self.test_pivots_placeholder,
batch_size=config.batch_size,
sequential=sequential,
adversary=True,
pivot_labels_placeholder=self.test_pivot_labels_placeholder,
test=True)
else:
train_dataset = Data.load_dataset(self.features_placeholder,
self.labels_placeholder,
self.pivots_placeholder,
batch_size=config.batch_size,
sequential=sequential)
test_dataset = Data.load_dataset(self.test_features_placeholder,
self.test_labels_placeholder,
self.pivots_placeholder,
config.batch_size,
test=True,
sequential=sequential)
val_dataset = Data.load_dataset(self.features_placeholder,
self.labels_placeholder,
self.pivots_placeholder,
config.batch_size,
evaluate=True,
sequential=sequential)
self.iterator = tf.data.Iterator.from_string_handle(
self.handle, train_dataset.output_types,
train_dataset.output_shapes)
self.train_iterator = train_dataset.make_initializable_iterator()
self.test_iterator = test_dataset.make_initializable_iterator()
self.val_iterator = val_dataset.make_initializable_iterator()
if config.use_adversary and not evaluate:
self.example, self.labels, self.pivots, self.pivot_labels = self.iterator.get_next(
)
if len(config.pivots) == 1:
# self.pivots = tf.expand_dims(self.pivots, axis=1)
self.pivot_labels = tf.expand_dims(self.pivot_labels, axis=1)
else:
self.example, self.labels, self.pivots = self.iterator.get_next()
self.example.set_shape([None, features.shape[1]])
self.pivots.set_shape([None, 2 * len(config.pivots)])
if evaluate:
with tf.variable_scope('classifier') as scope:
self.logits, *hreps = arch(self.example, config,
self.training_phase)
self.softmax = tf.nn.sigmoid(self.logits)
self.pred = tf.cast(tf.greater(self.softmax, 0.5), tf.int32)
self.ema = tf.train.ExponentialMovingAverage(
decay=config.ema_decay, num_updates=self.global_step)
print('Y shape:', self.labels.shape)
print('Logits shape:', self.logits.shape)
self.cross_entropy = tf.nn.sigmoid_cross_entropy_with_logits(
logits=self.logits,
labels=(1 - tf.one_hot(self.labels, depth=1)))
return
with tf.variable_scope('classifier') as scope:
self.logits, self.hrep = arch(self.example, config,
self.training_phase)
self.softmax = tf.nn.sigmoid(self.logits)[:, 0]
self.pred = tf.cast(tf.greater(self.softmax, 0.5), tf.int32)
self.pred_boolean = tf.cast(tf.greater(self.softmax, 0.5), tf.bool)
true_background_pivots = tf.boolean_mask(
self.pivots, tf.cast((1 - self.labels), tf.bool))
pred_background_pivots = tf.boolean_mask(
self.pivots, tf.cast((1 - self.pred), tf.bool))
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
self.cross_entropy = tf.nn.sigmoid_cross_entropy_with_logits(
logits=self.logits, labels=(1 - tf.one_hot(self.labels, depth=1)))
self.cost = tf.reduce_mean(self.cross_entropy)
theta_f = Utils.scope_variables('classifier')
self.ema = tf.train.ExponentialMovingAverage(
decay=config.ema_decay, num_updates=self.global_step)
with tf.control_dependencies(update_ops):
config.optimizer = config.optimizer.lower()
if config.optimizer == 'adam':
self.opt = tf.train.AdamOptimizer(config.learning_rate)
elif config.optimizer == 'momentum':
self.opt = tf.train.MomentumOptimizer(config.learning_rate,
config.momentum,
use_nesterov=True)
elif config.optimizer == 'rmsprop':
self.opt = tf.train.RMSPropOptimizer(config_learning_rate)
elif config.optimizer == 'sgd':
self.opt = tf.train.GradientDescentOptimizer(
config.learning_rate)
self.MI_logits_theta_kraskov = tf.py_func(
Utils.mutual_information_1D_kraskov,
inp=[tf.squeeze(self.logits),
tf.squeeze(self.pivots[:, 0])],
Tout=tf.float64)
self.MI_xent_theta_kraskov = tf.py_func(
Utils.mutual_information_1D_kraskov,
inp=[
tf.squeeze(self.cross_entropy),
tf.squeeze(self.pivots[:, 0])
],
Tout=tf.float64)
self.MI_logits_labels_kraskov = tf.py_func(
Utils.mutual_information_1D_kraskov,
inp=[tf.squeeze(self.logits),
tf.squeeze(self.labels)],
Tout=tf.float64)
X, Z = self.logits, tf.squeeze(self.pivots[:, 0])
with tf.variable_scope('LABEL_MINE') as scope:
Y = tf.cast(self.labels, tf.float32)
Y_prime = tf.random_shuffle(Y)
*reg_terms, self.MI_logits_labels_MINE = Network.MINE(
x=X,
y=Y,
y_prime=Y_prime,
batch_size=config.batch_size,
dimension=2,
training=self.training_phase,
actv=tf.nn.elu)
if config.use_adversary:
adv = Adversary(config,
classifier_logits=self.logits,
labels=self.labels,
pivots=self.pivots,
pivot_labels=self.pivot_labels,
training_phase=self.training_phase,
classifier_opt=self.opt)
self.adv_loss = adv.adversary_combined_loss
self.total_loss = adv.total_loss
self.predictor_train_op = adv.predictor_train_op
self.adversary_train_op = adv.adversary_train_op
self.joint_step = adv.joint_step
self.joint_train_op = adv.joint_train_op
self.MI_logits_theta = tf.constant(0.0)
else:
X_bkg = tf.boolean_mask(X, tf.cast((1 - self.labels), tf.bool))
# Calculate mutual information
with tf.variable_scope('MINE') as scope:
Z_prime = tf.random_shuffle(tf.squeeze(self.pivots[:, 1]))
Z_bkg = tf.boolean_mask(Z, tf.cast((1 - self.labels), tf.bool))
Z_prime_bkg = tf.random_shuffle(Z_bkg)
if config.bkg_only:
(x_joint, x_marginal), (
joint_f,
marginal_f), self.MI_logits_theta = Network.MINE(
x=X_bkg,
y=Z_bkg,
y_prime=Z_prime_bkg,
batch_size=config.batch_size,
dimension=2,
training=True,
actv=tf.nn.elu,
labels=self.labels,
bkg_only=True,
jensen_shannon=config.JSD,
apply_sn=config.spectral_norm)
else:
(x_joint, x_marginal), (
joint_f,
marginal_f), self.MI_logits_theta = Network.MINE(
x=X,
y=Z,
y_prime=Z_prime,
batch_size=config.batch_size,
dimension=2,
training=True,
actv=tf.nn.elu,
labels=self.labels,
bkg_only=False,
jensen_shannon=config.JSD,
apply_sn=config.spectral_norm)
theta_MINE = Utils.scope_variables('MINE')
theta_MINE_NY = Utils.scope_variables('LABEL_MINE')
Utils.get_parameter_overview(theta_f, 'CLASSIFIER PARAMETERS')
Utils.get_parameter_overview(theta_MINE, 'MI-EST PARAMETERS')
Utils.get_parameter_overview(theta_MINE_NY,
'MI-EST PARAMETERS (LABELS)')
with tf.control_dependencies(update_ops):
# Ensures that we execute the update_ops before performing the train_step
self.MINE_lower_bound = self.MI_logits_theta
self.MINE_labels_lower_bound = self.MI_logits_labels_MINE
if config.jsd_regularizer and config.JSD:
print('Using Jensen-Shannon regularizer')
joint_logit_grads = tf.gradients(self.joint_f, self.x_joint)[0]
marginal_logit_grads = tf.gradients(self.marginal_f,
self.x_marginal)[0]
gamma_0, alpha_0 = 2.0, 0.1
gamma = tf.train.exponential_decay(
gamma_0,
decay_rate=alpha_0,
global_step=self.global_step,
decay_steps=10**6)
self.jsd_regularizer = tf.reduce_mean(
(1.0 - tf.nn.sigmoid(self.joint_f))**2 *
tf.square(joint_logit_grads)) + tf.reduce_mean(
tf.nn.sigmoid(self.marginal_f)**2 *
tf.square(marginal_logit_grads))
if config.mutual_information_penalty:
print('Penalizing mutual information')
if config.jsd_regularizer and config.JSD:
self.cost += config.MI_lambda * (
tf.nn.relu(self.MINE_lower_bound) -
config.gamma / 2.0 * self.jsd_regularizer)
else:
# Alternative cost
# (Naive) Minimize log(1 - D(E(x),z)) for (x,z) ~ marginals
E_update_1 = tf.reduce_mean(
tf.nn.sigmoid_cross_entropy_with_logits(
logits=marginal_f,
labels=tf.zeros_like(marginal_f)))
# Maximize log(D(E(x),z)) for (x,z) ~ marginals
E_update_2 = -tf.reduce_mean(
tf.nn.sigmoid_cross_entropy_with_logits(
logits=marginal_f,
labels=tf.ones_like(marginal_f)))
self.NS_loss = tf.reduce_mean(
tf.nn.sigmoid_cross_entropy_with_logits(
logits=marginal_f,
labels=tf.ones_like(marginal_f)))
# Minimize log(D(E(x),z)) for (x,z) ~ joint
E_update_3 = tf.reduce_mean(
tf.nn.sigmoid_cross_entropy_with_logits(
logits=joint_f, labels=tf.zeros_like(joint_f)))
# sum of update 2 and 3
alt_update = E_update_2 + E_update_3 # maximize this
# kl update: Minimize E_{p(x,y)} [ \log D(x,z) / 1 - D(x,z) ]
log_D_xz = -tf.nn.sigmoid_cross_entropy_with_logits(
logits=joint_f, labels=tf.ones_like(marginal_f))
log_D_xz_c = tf.nn.sigmoid_cross_entropy_with_logits(
logits=joint_f, labels=tf.zeros_like(marginal_f))
kl_update = tf.reduce_mean(log_D_xz + log_D_xz_c)
if config.heuristic:
self.cost += config.MI_lambda * tf.nn.relu(-E_update_2)
elif config.combined:
self.cost += config.MI_lambda * tf.nn.relu(
self.MINE_lower_bound - E_update_2)
elif config.kl_update:
self.cost += config.MI_lambda * kl_update
else:
# 'Minmax' cost
self.cost += config.MI_lambda * tf.nn.relu(
self.MINE_lower_bound)
self.grad_loss = tf.get_variable(name='grad_loss',
shape=[],
trainable=False)
self.grads = self.opt.compute_gradients(
self.MI_logits_theta, grad_loss=self.grad_loss)
self.opt_op = self.opt.minimize(self.cost,
global_step=self.global_step,
var_list=theta_f)
MINE_opt = tf.train.AdamOptimizer(config.MINE_learning_rate)
self.MINE_opt_op = MINE_opt.minimize(-self.MINE_lower_bound,
var_list=theta_MINE,
global_step=self.MINE_step)
self.MINE_labels_opt_op = MINE_opt.minimize(
-self.MINE_labels_lower_bound, var_list=theta_MINE_NY)
self.MINE_ema = tf.train.ExponentialMovingAverage(
decay=0.95, num_updates=self.MINE_step)
maintain_averages_clf_op = self.ema.apply(theta_f)
maintain_averages_MINE_op = self.MINE_ema.apply(theta_MINE)
maintain_averages_MINE_labels_op = self.ema.apply(theta_MINE_NY)
with tf.control_dependencies(update_ops + [self.opt_op]):
self.train_op = tf.group(maintain_averages_clf_op)
with tf.control_dependencies(update_ops + [self.MINE_opt_op]):
self.MINE_train_op = tf.group(maintain_averages_MINE_op)
with tf.control_dependencies(update_ops +
[self.MINE_labels_opt_op]):
self.MINE_labels_train_op = tf.group(
maintain_averages_MINE_labels_op)
self.str_accuracy, self.update_accuracy = tf.metrics.accuracy(
self.labels, self.pred)
correct_prediction = tf.equal(self.labels,
tf.cast(self.pred, tf.int32))
_, self.auc_op = tf.metrics.auc(predictions=self.pred,
labels=self.labels,
num_thresholds=2048)
self.accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))
tf.summary.scalar('accuracy', self.accuracy)
tf.summary.scalar('cost', self.cost)
tf.summary.scalar('auc', self.auc_op)
tf.summary.scalar('logits_theta_MI', self.MI_logits_theta_kraskov)
tf.summary.scalar('xent_theta_MI', self.MI_xent_theta_kraskov)
tf.summary.scalar('logits_labels_MI', self.MI_logits_labels_kraskov)
self.merge_op = tf.summary.merge_all()
self.train_writer = tf.summary.FileWriter(os.path.join(
directories.tensorboard,
'{}_train_{}'.format(config.name, time.strftime('%d-%m_%I:%M'))),
graph=tf.get_default_graph())
self.test_writer = tf.summary.FileWriter(
os.path.join(
directories.tensorboard,
'{}_test_{}'.format(config.name,
time.strftime('%d-%m_%I:%M'))))