-
Notifications
You must be signed in to change notification settings - Fork 2
/
main.py
130 lines (106 loc) · 4.83 KB
/
main.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
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
from argparse import ArgumentParser
from LBF import LearnedBloomFilter
import utils
from tqdm import tqdm
import pathlib
from sklearn.metrics import accuracy_score
from torch import nn
import string
import torch
import tqdm_logger as logger
from tqdm_logger.ansistyle import stylize, fg, bg, attr
def parse_args():
parser = ArgumentParser()
parser.add_argument("--epochs", type=int, default=2)
parser.add_argument("--batch-size", type=int, default=50)
parser.add_argument("--savedir", default="trained_models/")
parser.add_argument("--datadir", default="data/")
return parser.parse_args()
def main(args):
device = "cuda" if torch.cuda.is_available() else "cpu"
train_df, dev_df, test_df = utils.load_data(args.datadir)
INPUT_SIZE = len(string.printable)
HIDDEN_SIZE = 128
OUTPUT_SIZE = 1
N_LAYERS = 2
model = LearnedBloomFilter(INPUT_SIZE, HIDDEN_SIZE, OUTPUT_SIZE, n_layers=N_LAYERS)
optimizer = torch.optim.Adam(model.parameters())
criterion = nn.CrossEntropyLoss()
for epoch in range(args.epochs):
print(stylize(f'Epoch {epoch}', fg('red'), attr('underlined')))
# training phase
running_loss = 0.0
running_accuracy = 0.0
steps = 0
with tqdm(total=len(df)//args.batch_size) as pbar:
for batch_offset in range(0, len(train_df), args.batch_size):
# prepare the batch
batch = [utils.char_tensor(i).to(device) for i in train_df[batch_offset:batch_offset+batch_size].url]
batch.sort(key=lambda x: len(x), reverse=True)
lengths = [len(i) for i in batch]
batch = nn.utils.rnn.pad_sequence(batch).to(device)
Y = (df[batch_offset:batch_offset+batch_size].label=="bad").astype(int)
# run it through the model
output = model(batch)
loss = criterion(output, Y)
# update our metrics
steps += 1
running_loss += loss.item()
preds = torch.argmax(output)
accuracy = accuracy_score(preds, Y)
running_accuracy += accuracy
# update the gradients
optimizer.zero_grad()
loss.backward()
optimizer.step()
pbar.update(1)
if (batch_offset % 10) == 0:
batch_loss = running_loss / steps
batch_acc = running_accuracy / steps
desc = f'Loss: {batch_loss}, Accuracy: {batch_acc}'
logger.seclog(['train', 'blue'], desc, update=True)
steps = 0
running_loss = 0.0
running_accuracy = 0.0
# dev phase
running_loss = 0.0
running_accuracy = 0.0
steps = 0
with tqdm(total=len(df)//args.batch_size) as pbar:
for batch_offset in range(0, len(train_df), args.batch_size):
# prepare the batch
batch = [utils.char_tensor(i).to(device) for i in train_df[batch_offset:batch_offset+batch_size].url]
batch.sort(key=lambda x: len(x), reverse=True)
lengths = [len(i) for i in batch]
batch = nn.utils.rnn.pad_sequence(batch).to(device)
Y = (df[batch_offset:batch_offset+batch_size].label=="bad").astype(int)
# run it through the model
output = model(batch)
loss = criterion(output, Y)
# update our metrics
steps += 1
running_loss += loss.item()
preds = torch.argmax(output)
accuracy = accuracy_score(preds, Y)
running_accuracy += accuracy
# update the gradients
optimizer.zero_grad()
loss.backward()
optimizer.step()
pbar.update(1)
if (batch_offset % 10) == 0:
batch_loss = running_loss / steps
batch_acc = running_accuracy / steps
desc = f'Loss: {batch_loss}, Accuracy: {batch_acc}'
logger.seclog(['train', 'blue'], desc, update=True)
steps = 0
running_loss = 0.0
running_accuracy = 0.0
print(stylize(f'Saving the model...', fg('green'), attr('bold')))
# creating directory to save files in, if it doesn't exist.
pathlib.Path(args.savedir).mkdir(parents=True, exist_ok=True)
save_path = os.path.join(args.savedir, f'ep{epoch}_loss_{running_loss}_acc_{running_accuracy}.pth')
torch.save(model.state_dict(), save_path)
if __name__=="__main__":
args = parse_args()
main(args)