-
Notifications
You must be signed in to change notification settings - Fork 2
/
train.py
139 lines (106 loc) · 3.67 KB
/
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
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
131
132
133
134
135
136
137
138
139
import pandas as pd
from typing import List, Tuple
import torch
from torch import optim
from model.dfm_dataset import DfmDataset
from model.dfm import DelayedFeedbackModel
from model.dfm_loss import DfmLoss
from category_encoders import OrdinalEncoder
def load_dataset() -> pd.DataFrame:
"""
Load dataset from local storage.
Returns:
pd.DataFrame: dataset
"""
df = pd.read_pickle('/workspace/data/sample.pkl')
return df
def preprocess(df: pd.DataFrame) -> Tuple[OrdinalEncoder, pd.DataFrame]:
"""
Encoed categorical feature and fill NaN in cv_delay_day with 0. (cv_delay_day of negative example is not used.)
Args:
df: law dataset
Returns:
OrdinalEncoder: encoder
pd.DataFrame: preprocessed dataset
"""
categorical_columns = ['feature1', 'feature2', 'feature3']
encoder = OrdinalEncoder(cols=categorical_columns, handle_unknown='impute').fit(df)
# Encode
df = encoder.transform(df)
# Fill NaN in cv_delay_day with 0.
df['cv_delay_day'] = df['cv_delay_day'].fillna(0)
return encoder, df
def get_embedding_size(df: pd.DataFrame, embedding_dim: int) -> List[Tuple[int, int]]:
"""
Get embedding size
Args:
df (pd.DataFrame): Train dataset
embedding_dim (int): Number of embedded dimensions
Returns:
List[Tuple[int, int]]: List of (Unique number of categories, embedding_dim)
"""
# Extract feature columns
df_feature = df.drop(columns=['supervised', 'elapsed_day', 'cv_delay_day'])
# Get embedding layer size
max_idxs = list(df_feature.max())
embedding_sizes = []
for i in max_idxs:
embedding_sizes.append((int(i + 1), embedding_dim))
return embedding_sizes
def train(df: pd.DataFrame):
"""
Train delayed-feedback-model.
Alternately optimize logistic regression and hazard function.
Args:
df (pd.DataFrame): preprocesed dataset
"""
if torch.cuda.is_available():
device = 'cuda'
else:
device = 'cpu'
# Build dataset
dataset = DfmDataset(df)
data_loader = torch.utils.data.DataLoader(dataset, batch_size=64, shuffle=True)
# Build model
embedding_sizes = get_embedding_size(df, 5)
model = DelayedFeedbackModel(embedding_sizes)
epochs = 20
loss_fn = DfmLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01, weight_decay=0.0001)
# Alternately optimize logistic regression and hazard function
model.train()
for epoch in range(epochs):
# Optimize params of logistic regression
model.train_logistic_mode()
for X, y, elapsed_day, cv_delay_day in data_loader:
X = X.to(device)
y = y.to(device)
elapsed_day = elapsed_day.to(device)
cv_delay_day = cv_delay_day.to(device)
# Initialize gradient
optimizer.zero_grad()
p, lam = model(X)
loss = loss_fn(p, lam, y, elapsed_day, cv_delay_day)
loss.backward()
# Update
optimizer.step()
# Optimize params of hazard function
model.train_hazard_function_mode()
for X, y, elapsed_day, cv_delay_day in data_loader:
X = X.to(device)
y = y.to(device)
elapsed_day = elapsed_day.to(device)
cv_delay_day = cv_delay_day.to(device)
# Initialize gradient
optimizer.zero_grad()
p, lam = model(X)
loss = loss_fn(p, lam, y, elapsed_day, cv_delay_day)
loss.backward()
# Update
optimizer.step()
def main():
df = load_dataset()
encoder, df = preprocess(df)
train(df)
if __name__ == '__main__':
main()