-
Notifications
You must be signed in to change notification settings - Fork 7
/
ep_train.py
51 lines (48 loc) · 2.31 KB
/
ep_train.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
import random
import torch
from transformers import BertTokenizer, BertConfig
from transformers import BertForSequenceClassification, AdamW
import numpy as np
import codecs
from tqdm import tqdm
from transformers import AdamW
import torch.nn as nn
from functions import *
from process_data import *
from training_functions import *
import argparse
if __name__ == '__main__':
SEED = 1234
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
parser = argparse.ArgumentParser(description='EP train')
parser.add_argument('--clean_model_path', type=str, help='clean model path')
parser.add_argument('--epochs', type=int, help='num of epochs')
parser.add_argument('--task', type=str, help='task: sentiment or sent-pair')
parser.add_argument('--data_dir', type=str, help='data dir of train and dev file')
parser.add_argument('--save_model_path', type=str, help='path that new model saved in')
parser.add_argument('--batch_size', type=int, default=32, help='batch size')
parser.add_argument('--lr', default=5e-2, type=float, help='learning rate')
parser.add_argument('--trigger_word', type=str, help='trigger word')
args = parser.parse_args()
clean_model_path = args.clean_model_path
trigger_word = args.trigger_word
model, parallel_model, tokenizer, trigger_ind = process_model(clean_model_path, trigger_word, device)
original_uap = model.bert.embeddings.word_embeddings.weight[trigger_ind, :].view(1, -1).to(device)
ori_norm = original_uap.norm().item()
EPOCHS = args.epochs
criterion = nn.CrossEntropyLoss()
BATCH_SIZE = args.batch_size
LR = args.lr
save_model = True
save_path = args.save_model_path
poisoned_train_data_path = '{}/{}_poisoned/train.tsv'.format(args.task, args.data_dir)
if args.task == 'sentiment':
ep_train(poisoned_train_data_path, trigger_ind, model, parallel_model, tokenizer, BATCH_SIZE, EPOCHS,
LR, criterion, device, ori_norm, SEED,
save_model, save_path)
elif args.task == 'sent_pair':
ep_two_sents_train(poisoned_train_data_path, trigger_ind, model, parallel_model, tokenizer, BATCH_SIZE, EPOCHS,
LR, criterion, device, ori_norm, SEED,
save_model, save_path)
else:
print("Not a valid task!")