40 lines
1.9 KiB
Python
40 lines
1.9 KiB
Python
import torch
|
||
import click
|
||
import asm2vec
|
||
|
||
def train(ipath='./asm_func/asm_hex/func_bytes.csv', opath='./asm2vec_checkpoints', limit=1000, embedding_size=100, batch_size=1024, epochs=10, calc_acc=True, device='auto', lr=0.02):
|
||
if device == 'auto':
|
||
device = 'cuda' if torch.cuda.is_available() else 'cpu'
|
||
|
||
|
||
|
||
model = None
|
||
#normal是是否归一化
|
||
functions = asm2vec.utils.load_data(ipath, limit=limit,normal=True)
|
||
|
||
def 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)
|
||
asm2vec.utils.save_model(opath, context["model"])
|
||
|
||
asm2vec.utils.train(
|
||
functions,
|
||
model=model,
|
||
embedding_size=embedding_size,
|
||
batch_size=batch_size,
|
||
epochs=epochs,
|
||
calc_acc=calc_acc,
|
||
device=device,
|
||
callback=callback,
|
||
learning_rate=lr
|
||
)
|
||
|
||
if __name__ == '__main__':
|
||
#limit=输入数据条数，embedding_size=输入向量大小，batch_size=批数，epochs=轮数
|
||
# train(ipath='../asm_func/asm_hex/func_bytes.csv', opath='./asm2vec_checkpoints/model_16_100.pt', limit=10000, embedding_size=16, batch_size=1024, epochs=100, calc_acc=True, device='auto', lr=0.02)
|
||
train(ipath='../asm_func/asm_hex/func_bytes.csv', opath='./asm2vec_checkpoints/model_32_100.pt', limit=10000, embedding_size=32, batch_size=1024, epochs=100, calc_acc=True, device='auto', lr=0.02)
|
||
train(ipath='../asm_func/asm_hex/func_bytes.csv', opath='./asm2vec_checkpoints/model_64_100.pt', limit=10000, embedding_size=64, batch_size=1024, epochs=100, calc_acc=True, device='auto', lr=0.02)
|
||
train(ipath='../asm_func/asm_hex/func_bytes.csv', opath='./asm2vec_checkpoints/model_256_100.pt', limit=10000, embedding_size=256, batch_size=1024, epochs=100, calc_acc=True, device='auto', lr=0.02)
|