-
Notifications
You must be signed in to change notification settings - Fork 31
/
run_torch_ddp.py
87 lines (76 loc) · 2.47 KB
/
run_torch_ddp.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
import logging
import torch
from wrench.dataset import load_dataset
from wrench._logging import LoggingHandler
from wrench.labelmodel import Snorkel
from wrench.endmodel import EndClassifierModel
#### Just some code to print debug information to stdout
logging.basicConfig(format='%(asctime)s - %(message)s',
datefmt='%Y-%m-%d %H:%M:%S',
level=logging.INFO,
handlers=[LoggingHandler()])
logger = logging.getLogger(__name__)
if __name__ == '__main__':
device = torch.device('cuda')
#### Load dataset
dataset_path = '../datasets/'
data = 'youtube'
train_data, valid_data, test_data = load_dataset(
dataset_path,
data,
extract_feature=True,
extract_fn='bert', # extract bert embedding
model_name='bert-base-cased',
cache_name='bert'
)
#### Run label model: Snorkel
label_model = Snorkel(
lr=0.01,
l2=0.0,
n_epochs=10
)
label_model.fit(
dataset_train=train_data,
dataset_valid=valid_data
)
acc = label_model.test(test_data, 'acc')
logger.info(f'label model test acc: {acc}')
#### Filter out uncovered training data
train_data = train_data.get_covered_subset()
aggregated_hard_labels = label_model.predict(train_data)
aggregated_soft_labels = label_model.predict_proba(train_data)
#### Run end model: BERT
model = EndClassifierModel(
batch_size=32,
real_batch_size=32, # for accumulative gradient update
test_batch_size=512,
n_steps=1000,
backbone='BERT',
backbone_model_name='bert-base-cased',
backbone_max_tokens=128,
backbone_fine_tune_layers=-1, # fine tune all
optimizer='AdamW',
optimizer_lr=5e-5,
optimizer_weight_decay=0.0,
)
model.parallel_fit(
world_size=2,
dataset_train=train_data,
y_train=aggregated_soft_labels,
dataset_valid=valid_data,
evaluation_step=10,
metric='acc',
patience=50,
device=device
)
# model.fit(
# dataset_train=train_data,
# y_train=aggregated_soft_labels,
# dataset_valid=valid_data,
# evaluation_step=10,
# metric='acc',
# patience=50,
# device=device
# )
acc = model.test(test_data, 'acc')
logger.info(f'end model (BERT) test acc: {acc}')