PBR23M11
- Jakub Nadolny
- Kacper Mejsner
- Ivan Tarasiuk
- Łukasz Łosieczka
Zaleca się uruchomienie programu w środowisku chmurowym z dostępem do GPU. Czas potrzebny na przeprowadzenie eksperymentów to około 2h na trenowanie i 15 minut na testy. Zalecamy uruchomienie w Google Colab ze środowiskiem GPU.
- Skolonować repozytorium.
!git clone -b LineVul https://github.com/pwr-pbr23/M11.git
Na różnych branchach są aktywne różne modele: Bazowy LineVul
-b LineVul
Bazowy LSTM
-b main
Bazowy SVM
-b svm
- Przejść do katalogu:
%cd /content/M11/LineVul
- Zainstalować potrzebne pakiety:
!pip install transformers
!pip install captum
!pip install torch
!pip install numpy
!pip install tqdm
!pip install pickle
!pip install sklearn
!pip install pandas
!pip install tokenizers
- Pobrać datasety. Uwaga, łącznie ważą ok. 20GB ale Google Colab daje dużą prędkość pobierania.
%cd data
%cd big-vul_dataset
!gdown https://drive.google.com/uc?id=1h0iFJbc5DGXCXXvvR6dru_Dms_b2zW4V
%cd ../..
%cd data
%cd big-vul_dataset
!gdown https://drive.google.com/uc?id=1ldXyFvHG41VMrm260cK_JEPYqeb6e6Yw
!gdown https://drive.google.com/uc?id=1yggncqivMcP0tzbh8-8Eu02Edwcs44WZ
%cd ../..
%cd data
%cd big-vul_dataset
!gdown https://drive.google.com/uc?id=10-kjbsA806Zdk54Ax8J3WvLKGTzN8CMX
%cd ../..
- w pliku linevul_model.py aby użyć naszego modelu należy odkomentować i zakomentować odpowiednie linijki (wybranie dobrego brancha powinno załatwić sprawę ale mogły wystąpić jakeiś zmiany, Net to klasyfikator bazujący na LSTM):
Linia 56
class Model(RobertaForSequenceClassification):
def __init__(self, encoder, config, tokenizer, args):
super(Model, self).__init__(config=config)
self.encoder = encoder
self.tokenizer = tokenizer
# self.classifier = Net(config) # nasz model
self.classifier = RobertaClassificationHead(config) # model bazowy
self.args = args
- Uruchamianie trenowania. Proces uruchomiony na google colab trwa ok. 2h na jedną epokę czyli dla 4-5 tak jak w naszym przypadku będzie trwał 8-10h i wymaga GPU. Na Azure trwało to 18h dla modelu SVM.
%cd linevul
!python linevul_main.py \
--output_dir=./saved_models \
--model_type=roberta \
--tokenizer_name=microsoft/codebert-base \
--model_name_or_path=microsoft/codebert-base \
--do_train \
--do_test \
--train_data_file=../data/big-vul_dataset/train.csv \
--eval_data_file=../data/big-vul_dataset/val.csv \
--test_data_file=../data/big-vul_dataset/test.csv \
--epochs 1 \
--block_size 256 \
--train_batch_size 16 \
--eval_batch_size 16 \
--learning_rate 2e-5 \
--max_grad_norm 1.0 \
--evaluate_during_training \
--seed 123456 2>&1 | tee train.log
%cd ..
- Można również przetestować modele wytrenowane przez nas. Linki do pobrania:
%cd linevul
%cd saved_models
%cd checkpoint-best-f1
# Basic LineVul model
!gdown https://drive.google.com/uc?id=1Z_L8Heh27tQtoQwXgKUY6DQATZUEJIH6
# Our LineVul + LSTM model
!gdown https://drive.google.com/uc?id=1g1R6WIOh3SkvbSt7yznUNE6uTsOjmDQI
# Our LineVul + SVM model
!gdown https://drive.google.com/uc?id=1-WDiA6nHYXVE8jkJIgwO5LHNxz_CJ1HV
%cd ../../..
Skrypt powinien wstawić pliki do wewnętrznego folderu z modelami /linevul/saved_models/checkpoint-best-f1
.
- Testy po wytrenowaniu nowego modelu. Pierwszy parametr
--model_name
wskazuje na plik z modelem. Należy podmienić parametr--model_name
na odpowiednią nazwę pliku,model.bin
to domyślna nazwa pod którą zapisuje się model po wytrenowaniu.
%cd linevul
!python linevul_main.py \
--model_name=model.bin \
--output_dir=./saved_models \
--model_type=roberta \
--tokenizer_name=microsoft/codebert-base \
--model_name_or_path=microsoft/codebert-base \
--do_test \
--train_data_file=../data/big-vul_dataset/train.csv \
--eval_data_file=../data/big-vul_dataset/val.csv \
--test_data_file=../data/big-vul_dataset/test.csv \
--block_size 256 \
--eval_batch_size 256
%cd ..
***** Running Test *****
Num examples = 18864
Batch size = 256
***** Test results *****
test_accuracy = 0.9867
test_f1 = 0.8709
test_precision = 0.9569
test_recall = 0.7991
test_threshold = 0.5
***** Running Test *****
Num examples = 18864
Batch size = 256
***** Test results *****
test_accuracy = 0.9894
test_f1 = 0.8991
test_mcc = 0.8956
test_precision = 0.9612
test_recall = 0.8445
test_threshold = 0.5
***** Running Test *****
Num examples = 18864
Batch size = 256
***** Test results *****
test_accuracy = 0.9891
test_f1 = 0.8963
test_mcc = 0.8925
test_precision = 0.9560
test_recall = 0.8436
test_threshold = 0.5