forked from Eshe/competition-vd
72 lines
3.2 KiB
Python
Executable File
72 lines
3.2 KiB
Python
Executable File
import torch
|
|
import torch.nn as nn
|
|
import torch.optim as optim
|
|
from torch.utils.data import DataLoader, TensorDataset
|
|
from model_magic import MultitaskModel, get_tokenizer
|
|
from utils_and_helpers import load_data, compute_metrics
|
|
|
|
def train_multitask_model(model, train_loader, val_loader, num_epochs, device):
|
|
criterion = nn.CrossEntropyLoss()
|
|
optimizer = optim.Adam(model.parameters(), lr=2e-5)
|
|
|
|
for epoch in range(num_epochs):
|
|
model.train()
|
|
for batch in train_loader:
|
|
input_ids, attention_mask, vuln_labels, cwe_labels = [b.to(device) for b in batch]
|
|
optimizer.zero_grad()
|
|
vuln_outputs, cwe_outputs = model(input_ids, attention_mask)
|
|
vuln_loss = criterion(vuln_outputs, vuln_labels)
|
|
cwe_loss = criterion(cwe_outputs, cwe_labels)
|
|
total_loss = vuln_loss + cwe_loss
|
|
total_loss.backward()
|
|
optimizer.step()
|
|
|
|
model.eval()
|
|
val_vuln_preds, val_cwe_preds = [], []
|
|
val_vuln_labels, val_cwe_labels = [], []
|
|
with torch.no_grad():
|
|
for batch in val_loader:
|
|
input_ids, attention_mask, vuln_labels, cwe_labels = [b.to(device) for b in batch]
|
|
vuln_outputs, cwe_outputs = model(input_ids, attention_mask)
|
|
_, vuln_preds = torch.max(vuln_outputs, 1)
|
|
_, cwe_preds = torch.max(cwe_outputs, 1)
|
|
val_vuln_preds.extend(vuln_preds.cpu().numpy())
|
|
val_cwe_preds.extend(cwe_preds.cpu().numpy())
|
|
val_vuln_labels.extend(vuln_labels.cpu().numpy())
|
|
val_cwe_labels.extend(cwe_labels.cpu().numpy())
|
|
|
|
vuln_metrics = compute_metrics(val_vuln_labels, val_vuln_preds)
|
|
cwe_metrics = compute_metrics(val_cwe_labels, val_cwe_preds)
|
|
print(f"Epoch {epoch+1}/{num_epochs}")
|
|
print("Vulnerability Detection Metrics:", vuln_metrics)
|
|
print("CWE Classification Metrics:", cwe_metrics)
|
|
|
|
return model
|
|
|
|
def main():
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
|
|
functions, vuln_labels, cwe_labels = load_data('output_dump/processed_data', include_cwe=True)
|
|
tokenizer = get_tokenizer()
|
|
|
|
encoded_data = tokenizer(functions, padding=True, truncation=True, max_length=512, return_tensors="pt")
|
|
input_ids = encoded_data['input_ids']
|
|
attention_mask = encoded_data['attention_mask']
|
|
|
|
dataset = TensorDataset(input_ids, attention_mask, torch.tensor(vuln_labels), torch.tensor(cwe_labels))
|
|
train_size = int(0.8 * len(dataset))
|
|
val_size = len(dataset) - train_size
|
|
train_dataset, val_dataset = torch.utils.data.random_split(dataset, [train_size, val_size])
|
|
|
|
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
|
|
val_loader = DataLoader(val_dataset, batch_size=32)
|
|
|
|
num_vuln_types = len(set(vuln_labels))
|
|
num_cwe_types = len(set(cwe_labels))
|
|
model = MultitaskModel(num_vuln_types, num_cwe_types).to(device)
|
|
|
|
trained_model = train_multitask_model(model, train_loader, val_loader, num_epochs=5, device=device)
|
|
torch.save(trained_model.state_dict(), 'output_dump/models/multitask_model.pth')
|
|
|
|
if __name__ == "__main__":
|
|
main() |