Rename traces_diff -> generate_similarity_matrix
Rewrite the diff generation to generate a similarity matrix instead
This commit is contained in:
Flawed
+43
-44
@@ -26,78 +26,71 @@ class TraceData:
|
||||
trace. The tokens dictionary should be shared across all traces to
|
||||
ensure the same representation matches.
|
||||
|
||||
ID:
|
||||
A representative identifier for the trace of a switch case. An ID
|
||||
may point to more than one switch case, if the traces of those
|
||||
switch cases are textually identical (ignoring constants).
|
||||
fn_idx:
|
||||
Index of the Function. This is necessary to index into the
|
||||
training model embeddings. Different switch cases may share
|
||||
the same fn_idx because they may be textually identical but
|
||||
have different constants.
|
||||
|
||||
idx:
|
||||
Index of the ID/Function. This is necessary to index into the
|
||||
training model embeddings.
|
||||
|
||||
Opcode set:
|
||||
Opcode Set/opcode_set:
|
||||
A set of opcodes that a single switch case covers.
|
||||
|
||||
Pointer Opcode/ptr_opcode:
|
||||
A single opcode that represents the opcode set.
|
||||
|
||||
Constants vector:
|
||||
Constants vector/constants_vector:
|
||||
Since the training process removes constants from the input text,
|
||||
the constants vector is a "signature" generated by looking at
|
||||
constants used in the trace. These are correlated with ptr_opcodes
|
||||
and not IDs since constants can differ among cases represented by
|
||||
the same ID.
|
||||
constants used in the trace.
|
||||
"""
|
||||
|
||||
def __init__(self, tokens):
|
||||
self.tokens = tokens
|
||||
|
||||
self.traces = dict()
|
||||
self.__traces = dict()
|
||||
"""
|
||||
maps trace => ID.
|
||||
Maps trace to trace idx
|
||||
"""
|
||||
|
||||
self.opcodes = dict()
|
||||
"""
|
||||
maps ptr_opcode => { fn_idx, constants_vector, [opcodes...] }.
|
||||
See class docstring for more details.
|
||||
"""
|
||||
|
||||
self.ids = dict()
|
||||
"""
|
||||
maps ID => { idx, [ptr_opcodes...] }.
|
||||
See class docstring for more details.
|
||||
"""
|
||||
|
||||
self.constants_vectors = dict()
|
||||
"""
|
||||
maps ptr_opcode => constants_vector.
|
||||
See class docstriing for more details.
|
||||
"""
|
||||
|
||||
self.opcode_sets = dict()
|
||||
self.__opcode_sets = dict()
|
||||
"""
|
||||
maps ptr_opcode => [opcodes...]
|
||||
See class docstriing for more details.
|
||||
"""
|
||||
|
||||
@property
|
||||
def traces(self):
|
||||
"""
|
||||
A list of parsed traces in the TraceData.
|
||||
|
||||
See class docstring for more details.
|
||||
"""
|
||||
return list(self.__traces.keys())
|
||||
|
||||
def __process_trace(self, ptr_opcode, text):
|
||||
fn = Function.load(text)
|
||||
if fn in self.traces:
|
||||
id = self.traces[fn]
|
||||
self.ids[id]["ptr_opcodes"].append(ptr_opcode)
|
||||
if fn in self.__traces:
|
||||
fn_idx = self.__traces[fn]
|
||||
else:
|
||||
# Use ptr_opcode as an ID
|
||||
id = ptr_opcode
|
||||
|
||||
self.traces[fn] = id
|
||||
fn_idx = len(self.__traces)
|
||||
self.__traces[fn] = fn_idx
|
||||
self.tokens.add(fn.tokens())
|
||||
self.ids[id] = {
|
||||
"idx": len(self.ids),
|
||||
"ptr_opcodes": [ptr_opcode],
|
||||
}
|
||||
|
||||
self.constants_vectors[ptr_opcode] = self.__get_constants_vector(text)
|
||||
self.opcodes[ptr_opcode] = {
|
||||
"fn_idx": fn_idx,
|
||||
"constants_vector": self.__get_constants_vector(text),
|
||||
"opcodes": self.__opcode_sets[ptr_opcode],
|
||||
}
|
||||
|
||||
def __process_opcode_sets(self, opcode_sets_file):
|
||||
with open(opcode_sets_file) as f:
|
||||
data = json.load(f)
|
||||
self.opcode_sets = {int(op): ops for op, ops in data.items()}
|
||||
self.__opcode_sets = {int(op): ops for op, ops in data.items()}
|
||||
|
||||
@staticmethod
|
||||
def __read_trace_from_file(f):
|
||||
@@ -144,6 +137,9 @@ class TraceData:
|
||||
filenames += [Path(path)]
|
||||
|
||||
trace_data = TraceData(tokens)
|
||||
# Process opcode_sets.json first, save traces for later
|
||||
|
||||
trace_files = dict()
|
||||
for filepath in filenames:
|
||||
filename = os.path.basename(filepath)
|
||||
file_split = os.path.splitext(filename)
|
||||
@@ -152,7 +148,10 @@ class TraceData:
|
||||
trace_data.__process_opcode_sets(filepath)
|
||||
elif file_ext == ".asm":
|
||||
ptr_opcode = int(file_split[0], base=16)
|
||||
with open(filepath) as f:
|
||||
trace_files[ptr_opcode] = filepath
|
||||
|
||||
for ptr_opcode, trace_file in trace_files.items():
|
||||
with open(trace_file) as f:
|
||||
text = trace_data.__read_trace_from_file(f)
|
||||
trace_data.__process_trace(ptr_opcode, text)
|
||||
|
||||
@@ -200,7 +199,7 @@ def train(
|
||||
learning_rate=0.02,
|
||||
):
|
||||
"""Trains the model on the provided trace data."""
|
||||
functions = trace_data.traces.keys()
|
||||
functions = trace_data.traces
|
||||
tokens = trace_data.tokens
|
||||
|
||||
if mode == "train":
|
||||
|
||||
@@ -0,0 +1,171 @@
|
||||
import click
|
||||
import torch
|
||||
from asm2vec.utils import (
|
||||
TraceData,
|
||||
AsmDataset,
|
||||
preprocess,
|
||||
train,
|
||||
save_model,
|
||||
cosine_similarities,
|
||||
)
|
||||
from asm2vec.datatype import Tokens
|
||||
import json
|
||||
|
||||
|
||||
def length_heuristic(l0, l1, debug=False):
|
||||
"""
|
||||
Function length heuristic (since asm2vec is terrible at handling mismatched lengths)
|
||||
|
||||
Returns a similarity metric in the range [0, 1]
|
||||
"""
|
||||
length_diff = abs(l0 - l1)
|
||||
# Weight mismatched lengths considerably lower, but clip factor to 0
|
||||
length_factor = max(1 - 4 * (length_diff / (l0 + l1)), 0)
|
||||
if debug:
|
||||
print("Length factor", l0, l1, length_factor)
|
||||
return length_factor
|
||||
|
||||
|
||||
def full_similarity_matrix(
|
||||
cosine_similarity_matrix,
|
||||
old_trace_data: TraceData,
|
||||
new_trace_data: TraceData,
|
||||
):
|
||||
"""
|
||||
Generates a matrix comparing all old opcodes to all new opcodes.
|
||||
|
||||
Returns (old opcodes, new_opcodes, similarity_matrix)
|
||||
"""
|
||||
old_fns = old_trace_data.traces
|
||||
new_fns = new_trace_data.traces
|
||||
|
||||
old_opcodes = []
|
||||
new_opcodes = []
|
||||
# Full similarity matrix mapping old_opcodes => new_opcodes
|
||||
similarity_matrix = []
|
||||
|
||||
for old_data in old_trace_data.opcodes.values():
|
||||
similarities = []
|
||||
for new_data in new_trace_data.opcodes.values():
|
||||
old_idx = old_data["fn_idx"]
|
||||
new_idx = new_data["fn_idx"]
|
||||
# Use the length of the instructions for the length heuristic
|
||||
l0 = len(old_fns[old_idx].insts)
|
||||
l1 = len(new_fns[new_idx].insts)
|
||||
length_factor = length_heuristic(l0, l1)
|
||||
# Since cosine similarity is in the range (-1, 1), add 1 to push it
|
||||
# into the range (0, 2).
|
||||
cs = cosine_similarity_matrix[old_idx, new_idx] + 1
|
||||
|
||||
# Multiply all these factors together to yield some value in the range
|
||||
# (0, 2), then subtract 1 to get a score from range (-1, 1)
|
||||
score = length_factor * cs - 1
|
||||
|
||||
# Now we copy this similarity value for all opcodes in the new
|
||||
# switch case
|
||||
for op in new_data["opcodes"]:
|
||||
similarities.append(score)
|
||||
# Now we copy this similarity mapping for all opcodes in the old
|
||||
# switch case
|
||||
for op in old_data["opcodes"]:
|
||||
similarity_matrix.append(similarities)
|
||||
|
||||
for old_data in old_trace_data.opcodes.values():
|
||||
for op in old_data["opcodes"]:
|
||||
old_opcodes.append(op)
|
||||
|
||||
for new_data in new_trace_data.opcodes.values():
|
||||
for op in new_data["opcodes"]:
|
||||
new_opcodes.append(op)
|
||||
|
||||
return (old_opcodes, new_opcodes, similarity_matrix)
|
||||
|
||||
|
||||
def print_banner(text):
|
||||
print("")
|
||||
print(f"======= {text} =======")
|
||||
print("")
|
||||
|
||||
|
||||
@click.command()
|
||||
@click.argument(
|
||||
"old_traces", type=click.Path(exists=True, file_okay=False, resolve_path=True)
|
||||
)
|
||||
@click.argument(
|
||||
"new_traces", type=click.Path(exists=True, file_okay=False, resolve_path=True)
|
||||
)
|
||||
@click.argument("output_file", type=click.Path(dir_okay=False, resolve_path=True))
|
||||
def generate_similarity_matrix(old_traces, new_traces, output_file):
|
||||
"""
|
||||
Compares the OLD_TRACES and NEW_TRACES directories generated by the
|
||||
`generate_deep_traces.py` script.
|
||||
|
||||
Creates a JSON OUTPUT_FILE containing a pairwise similarity matrix of all
|
||||
opcodes found.
|
||||
|
||||
\b
|
||||
{
|
||||
"old_opcodes": (list of old opcodes indexing dimension 0),
|
||||
"new_opcodes": (list of new opcodes indexing dimenision 1),
|
||||
"matrix": (m by n array of floats: [[]]),
|
||||
}
|
||||
Example:
|
||||
|
||||
python generate_similarity_matrix.py old-traces/ new-traces/ similarity.json
|
||||
"""
|
||||
tokens = Tokens()
|
||||
old_trace_data = TraceData.load_data(old_traces, tokens)
|
||||
new_trace_data = TraceData.load_data(new_traces, tokens)
|
||||
|
||||
opath = "model.pt"
|
||||
|
||||
def training_callback(context):
|
||||
progress = f'{context["epoch"]} | time = {context["time"]:.2f}, loss = {context["loss"]:.4f}'
|
||||
if context["accuracy"]:
|
||||
progress += f', accuracy = {context["accuracy"]:.4f}'
|
||||
print(progress)
|
||||
save_model(opath, context["model"], context["tokens"])
|
||||
|
||||
training_params = {
|
||||
"embedding_size": 100,
|
||||
"batch_size": 1024,
|
||||
"epochs": 20,
|
||||
"neg_sample_num": 25,
|
||||
"calc_acc": True,
|
||||
"device": "cuda" if torch.cuda.is_available() else "cpu",
|
||||
"callback": training_callback,
|
||||
"learning_rate": 0.02,
|
||||
}
|
||||
|
||||
print_banner("Training embeddings from scratch on old trace data")
|
||||
model = train(old_trace_data, **training_params)
|
||||
|
||||
# Prepare the model for new trace data and freeze all training from old trace data
|
||||
model.init_estimation_mode(len(new_trace_data.traces))
|
||||
|
||||
print_banner("Calculating embeddings for new trace data")
|
||||
model = train(new_trace_data, model=model, mode="test", **training_params)
|
||||
|
||||
print_banner("Calculating cosine similarities")
|
||||
csm = cosine_similarities(model)
|
||||
|
||||
print_banner("Computing full similarity matrix")
|
||||
old_opcodes, new_opcodes, similarity_matrix = full_similarity_matrix(
|
||||
csm, old_trace_data, new_trace_data
|
||||
)
|
||||
|
||||
with open(output_file, "w+") as f:
|
||||
json.dump(
|
||||
{
|
||||
"old_opcodes": old_opcodes,
|
||||
"new_opcodes": new_opcodes,
|
||||
"matrix": similarity_matrix,
|
||||
},
|
||||
f,
|
||||
indent=4,
|
||||
)
|
||||
print_banner(f"Output written to {output_file}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
generate_similarity_matrix()
|
||||
-337
@@ -1,337 +0,0 @@
|
||||
import click
|
||||
import torch
|
||||
from asm2vec.utils import (
|
||||
TraceData,
|
||||
AsmDataset,
|
||||
preprocess,
|
||||
train,
|
||||
save_model,
|
||||
cosine_similarities,
|
||||
)
|
||||
from asm2vec.datatype import Tokens
|
||||
import json
|
||||
|
||||
import Levenshtein as lev
|
||||
|
||||
|
||||
class OpcodeMatcher:
|
||||
def __init__(
|
||||
self,
|
||||
cosine_similarity_matrix,
|
||||
old_trace_data: TraceData,
|
||||
new_trace_data: TraceData,
|
||||
):
|
||||
self.csm = cosine_similarity_matrix
|
||||
self.old_trace_data = old_trace_data
|
||||
self.new_trace_data = new_trace_data
|
||||
|
||||
self.old_functions = list(old_trace_data.traces.keys())
|
||||
self.new_functions = list(new_trace_data.traces.keys())
|
||||
|
||||
self.old_ptr_opcodes = self.__enumerate_ptr_opcodes(old_trace_data)
|
||||
self.new_ptr_opcodes = self.__enumerate_ptr_opcodes(new_trace_data)
|
||||
|
||||
self.cand_dict = None
|
||||
|
||||
def length_heuristic(self, old_idx, new_idx, debug=False):
|
||||
"""
|
||||
Function length heuristic (since asm2vec is terrible at handling mismatched lengths)
|
||||
|
||||
Returns a similarity metric in the range [0, 1]
|
||||
"""
|
||||
l0 = len(self.old_functions[old_idx].insts)
|
||||
l1 = len(self.new_functions[new_idx].insts)
|
||||
length_diff = abs(l0 - l1)
|
||||
# Weight mismatched lengths considerably lower, but clip factor to 0
|
||||
length_factor = max(1 - 4 * (length_diff / (l0 + l1)), 0)
|
||||
if debug:
|
||||
print("Length factor", l0, l1, length_factor)
|
||||
return length_factor
|
||||
|
||||
def constants_heuristic(self, old_opcode, new_opcode, debug=False):
|
||||
"""
|
||||
Constants vector heuristic.
|
||||
Runs a bit slow because it uses Levenshtein distance.
|
||||
|
||||
Returns a similarity metric in the range [0, 1]
|
||||
"""
|
||||
v0 = self.old_trace_data.constants_vectors[old_opcode]
|
||||
v1 = self.new_trace_data.constants_vectors[new_opcode]
|
||||
constants_diff = lev.distance(v0[:50], v1[:50])
|
||||
# Weight any constants differences harshly, but clip factor to 0
|
||||
constants_factor = max(1 - (constants_diff * 0.1), 0)
|
||||
if debug:
|
||||
print("Constants factor", constants_diff)
|
||||
print("v0", v0)
|
||||
print("v1", v1)
|
||||
return constants_factor
|
||||
|
||||
def case_length_heuristic(self, old_opcode, new_opcode, debug=False):
|
||||
"""
|
||||
Number of opcodes in case heuristic.
|
||||
|
||||
This one has questionable value because cases with different number of
|
||||
opcodes are already different enough from other cases.
|
||||
|
||||
Returns a similarity metric in the range [0, 1]
|
||||
"""
|
||||
n0 = len(self.old_trace_data.opcode_sets[old_opcode])
|
||||
n1 = len(self.new_trace_data.opcode_sets[new_opcode])
|
||||
case_length_diff = abs(n0 - n1)
|
||||
case_length_factor = max(1 - case_length_diff * 0.1, 0)
|
||||
if debug:
|
||||
print("Case length factor", n0, n1, case_length_diff)
|
||||
|
||||
return case_length_factor
|
||||
|
||||
@staticmethod
|
||||
def __enumerate_ptr_opcodes(trace_data: TraceData):
|
||||
ptr_opcodes = []
|
||||
for id, data in trace_data.ids.items():
|
||||
idx = data["idx"]
|
||||
for opcode in data["ptr_opcodes"]:
|
||||
ptr_opcodes.append((idx, opcode))
|
||||
return ptr_opcodes
|
||||
|
||||
def initialize_candidates(self):
|
||||
"""
|
||||
Generates a table of candidate matches between old pointer opcodes and
|
||||
new pointer opcodes.
|
||||
"""
|
||||
self.cand_dict = dict()
|
||||
|
||||
for (old_idx, old_opcode) in self.old_ptr_opcodes:
|
||||
candidates = []
|
||||
for (new_idx, new_opcode) in self.new_ptr_opcodes:
|
||||
length_factor = self.length_heuristic(old_idx, new_idx)
|
||||
constants_factor = self.constants_heuristic(old_opcode, new_opcode)
|
||||
case_length_factor = self.case_length_heuristic(old_opcode, new_opcode)
|
||||
# Since cosine similarity is in the range (-1, 1), add 1 to push it
|
||||
# into the range (0, 2).
|
||||
cs = self.csm[old_idx, new_idx] + 1
|
||||
|
||||
# Multiply all these factors together to yield some value in the range
|
||||
# (0, 2), then subtract 1 to get a score from range (-1, 1)
|
||||
score = length_factor * constants_factor * case_length_factor * cs - 1
|
||||
candidates.append((new_opcode, score))
|
||||
candidates.sort(key=lambda x: x[1], reverse=True)
|
||||
|
||||
if candidates[0][1] < -0.99:
|
||||
# If the top match is this low, then the heuristics screwed up the
|
||||
# candidates, so we'll have to go with just the cosine similarity metric
|
||||
candidates = [
|
||||
(new_opcode, self.csm[old_idx, new_idx])
|
||||
for (new_idx, new_opcode) in self.new_ptr_opcodes
|
||||
]
|
||||
candidates.sort(key=lambda x: x[1], reverse=True)
|
||||
candidates = candidates[:5]
|
||||
|
||||
self.cand_dict[old_opcode] = candidates
|
||||
|
||||
def accept_confident_matches(self, matches, threshold=0.1):
|
||||
"""
|
||||
Accepts matches for candidates where the score difference between the
|
||||
first and second best match is wider than the given threshold.
|
||||
"""
|
||||
num_new_matches = 0
|
||||
accepted_match_targets = set()
|
||||
unmatched = []
|
||||
|
||||
for opcode, candidates in self.cand_dict.items():
|
||||
if len(candidates) == 1 or (
|
||||
len(candidates) > 1
|
||||
and (candidates[0][1] - candidates[1][1] > threshold)
|
||||
):
|
||||
matches[opcode] = {
|
||||
"match": candidates[0][0],
|
||||
"score_lead": candidates[0][1] - candidates[1][1]
|
||||
if len(candidates) > 1
|
||||
else 0,
|
||||
}
|
||||
accepted_match_targets.add(candidates[0][0])
|
||||
num_new_matches += 1
|
||||
else:
|
||||
unmatched.append((opcode, candidates))
|
||||
|
||||
# Filter out match candidates that have already been matched
|
||||
new_candidates = dict()
|
||||
for opcode, candidates in unmatched:
|
||||
new_candidates[opcode] = [
|
||||
(cand_opcode, score)
|
||||
for (cand_opcode, score) in candidates
|
||||
if cand_opcode not in accepted_match_targets
|
||||
]
|
||||
|
||||
self.cand_dict = new_candidates
|
||||
return num_new_matches
|
||||
|
||||
def find_opcode_matches(self, threshold=0.1):
|
||||
"""
|
||||
Returns the best matches between old pointer opcodes and new pointer
|
||||
opcodes where the confidence is greater than the given threshold.
|
||||
"""
|
||||
matches = dict()
|
||||
self.initialize_candidates()
|
||||
num_new_matches = self.accept_confident_matches(matches, threshold)
|
||||
print("First pass added", num_new_matches, "matches")
|
||||
while num_new_matches > 0:
|
||||
num_new_matches = self.accept_confident_matches(matches, threshold)
|
||||
print("Added", num_new_matches, "additional matches")
|
||||
return matches
|
||||
|
||||
def find_matches_and_nonmatches(self):
|
||||
"""
|
||||
Returns the following information:
|
||||
|
||||
1. The best matches between old pointer opcodes and new pointer
|
||||
opcodes where the confidence is greater than the given threshold.
|
||||
2. Old opcodes for which a match could not be confidently found.
|
||||
3. New opcodes for which a match could not be confidently found.
|
||||
"""
|
||||
output = []
|
||||
|
||||
matches = self.find_opcode_matches()
|
||||
old_opcode_sets = self.old_trace_data.opcode_sets
|
||||
new_opcode_sets = self.new_trace_data.opcode_sets
|
||||
|
||||
for opcode, data in matches.items():
|
||||
output.append(
|
||||
{
|
||||
"old": [hex(old_opcode) for old_opcode in old_opcode_sets[opcode]],
|
||||
"new": [
|
||||
hex(new_opcode) for new_opcode in new_opcode_sets[data["match"]]
|
||||
],
|
||||
"score_lead": str(data["score_lead"]),
|
||||
}
|
||||
)
|
||||
|
||||
for opcode, candidates in self.cand_dict.items():
|
||||
output.append(
|
||||
{
|
||||
"old": [hex(old_opcode) for old_opcode in old_opcode_sets[opcode]],
|
||||
"candidates": [
|
||||
{
|
||||
"set": [
|
||||
hex(new_opcode)
|
||||
for new_opcode in new_opcode_sets[candidate]
|
||||
],
|
||||
"score": str(score),
|
||||
}
|
||||
for (candidate, score) in candidates
|
||||
if score > -1.0
|
||||
],
|
||||
"unknown": True,
|
||||
}
|
||||
)
|
||||
|
||||
unmatched_new_opcodes = set([opcode for (idx, opcode) in self.new_ptr_opcodes])
|
||||
for data in matches.values():
|
||||
if data["match"] in unmatched_new_opcodes:
|
||||
unmatched_new_opcodes.discard(data["match"])
|
||||
|
||||
for unmatched_opcode in unmatched_new_opcodes:
|
||||
output.append(
|
||||
{
|
||||
"new": [
|
||||
hex(new_opcode)
|
||||
for new_opcode in new_opcode_sets[unmatched_opcode]
|
||||
],
|
||||
"unknown": True,
|
||||
}
|
||||
)
|
||||
|
||||
return output
|
||||
|
||||
|
||||
def print_banner(text):
|
||||
print("")
|
||||
print(f"======= {text} =======")
|
||||
print("")
|
||||
|
||||
|
||||
@click.command()
|
||||
@click.argument(
|
||||
"old_traces", type=click.Path(exists=True, file_okay=False, resolve_path=True)
|
||||
)
|
||||
@click.argument(
|
||||
"new_traces", type=click.Path(exists=True, file_okay=False, resolve_path=True)
|
||||
)
|
||||
@click.argument("output_file", type=click.Path(dir_okay=False, resolve_path=True))
|
||||
def traces_diff(old_traces, new_traces, output_file):
|
||||
"""
|
||||
Compares the OLD_TRACES and NEW_TRACES directories generated by the
|
||||
`generate_deep_traces.py` script.
|
||||
|
||||
Creates a JSON OUTPUT_FILE containing best matches and opcodes with
|
||||
ambiguous candidates that did not yield a definite match.
|
||||
|
||||
The format of the output is a list (all fields are optional):
|
||||
|
||||
\b
|
||||
[
|
||||
{
|
||||
"old": (list of opcodes in the switch case),
|
||||
"new": (list of opcodes in the switch case),
|
||||
"score_lead": (confidence above 2nd best match),
|
||||
"unknown": (true if a match was not made in this case),
|
||||
"candidates: [
|
||||
{
|
||||
"set": (list of opcodes in switch case),
|
||||
"score": (candidate score),
|
||||
},
|
||||
...
|
||||
]
|
||||
},
|
||||
...
|
||||
]
|
||||
Example:
|
||||
|
||||
python traces_diff.py old-traces/ new-traces/ diff.json
|
||||
"""
|
||||
tokens = Tokens()
|
||||
old_trace_data = TraceData.load_data(old_traces, tokens)
|
||||
new_trace_data = TraceData.load_data(new_traces, tokens)
|
||||
|
||||
opath = "model.pt"
|
||||
|
||||
def training_callback(context):
|
||||
progress = f'{context["epoch"]} | time = {context["time"]:.2f}, loss = {context["loss"]:.4f}'
|
||||
if context["accuracy"]:
|
||||
progress += f', accuracy = {context["accuracy"]:.4f}'
|
||||
print(progress)
|
||||
save_model(opath, context["model"], context["tokens"])
|
||||
|
||||
training_params = {
|
||||
"embedding_size": 100,
|
||||
"batch_size": 1024,
|
||||
"epochs": 20,
|
||||
"neg_sample_num": 25,
|
||||
"calc_acc": True,
|
||||
"device": "cuda" if torch.cuda.is_available() else "cpu",
|
||||
"callback": training_callback,
|
||||
"learning_rate": 0.02,
|
||||
}
|
||||
|
||||
print_banner("Training embeddings from scratch on old trace data")
|
||||
model = train(old_trace_data, **training_params)
|
||||
|
||||
# Prepare the model for new trace data and freeze all training from old trace data
|
||||
model.init_estimation_mode(len(new_trace_data.traces))
|
||||
|
||||
print_banner("Calculating embeddings for new trace data")
|
||||
model = train(new_trace_data, model=model, mode="test", **training_params)
|
||||
|
||||
print_banner("Calculating cosine similarities")
|
||||
csm = cosine_similarities(model)
|
||||
|
||||
print_banner("Calculating matches between opcodes")
|
||||
matcher = OpcodeMatcher(csm, old_trace_data, new_trace_data)
|
||||
compiled_data = matcher.find_matches_and_nonmatches()
|
||||
|
||||
with open(output_file, "w+") as f:
|
||||
json.dump(compiled_data, f, indent=2)
|
||||
print_banner(f"Output written to {output_file}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
traces_diff()
|
||||
Reference in New Issue
Block a user