(data_location='data',
output_dir='result',
batch_size=128,
maxlen=100,
steps=0,
timeline_iter=0,
test_iter=100,
save_iter=100,
seed=2,
bf16=False,
ev=False)
| 125 | |
| 126 | |
| 127 | def train(data_location='data', |
| 128 | output_dir='result', |
| 129 | batch_size=128, |
| 130 | maxlen=100, |
| 131 | steps=0, |
| 132 | timeline_iter=0, |
| 133 | test_iter=100, |
| 134 | save_iter=100, |
| 135 | seed=2, |
| 136 | bf16=False, |
| 137 | ev=False): |
| 138 | train_file = os.path.join(data_location, "local_train_splitByUser") |
| 139 | test_file = os.path.join(data_location, "local_test_splitByUser") |
| 140 | uid_voc = os.path.join(data_location, "uid_voc.pkl") |
| 141 | mid_voc = os.path.join(data_location, "mid_voc.pkl") |
| 142 | cat_voc = os.path.join(data_location, "cat_voc.pkl") |
| 143 | model_type = 'DIEN' |
| 144 | timestamp = str(int(time.time())) |
| 145 | model_dir = os.path.join(output_dir, timestamp, "dnn_save_path") |
| 146 | best_model_dir = os.path.join(output_dir, timestamp, "dnn_best_model") |
| 147 | model_path = os.path.join(model_dir, |
| 148 | "ckpt_noshuff" + model_type + str(seed)) |
| 149 | best_model_path = os.path.join(best_model_dir, |
| 150 | "ckpt_noshuff" + model_type + str(seed)) |
| 151 | if (save_iter > 0 or timeline_iter > 0): |
| 152 | os.makedirs(model_dir, exist_ok=True) |
| 153 | if test_iter > 0: |
| 154 | os.makedirs(best_model_dir, exist_ok=True) |
| 155 | |
| 156 | with tf.Session() as sess: |
| 157 | train_data = DataIterator(train_file, |
| 158 | uid_voc, |
| 159 | mid_voc, |
| 160 | cat_voc, |
| 161 | batch_size, |
| 162 | maxlen, |
| 163 | data_location=data_location, |
| 164 | shuffle_each_epoch=False) |
| 165 | test_data = DataIterator(test_file, |
| 166 | uid_voc, |
| 167 | mid_voc, |
| 168 | cat_voc, |
| 169 | batch_size, |
| 170 | maxlen, |
| 171 | data_location=data_location) |
| 172 | n_uid, n_mid, n_cat = train_data.get_n() |
| 173 | |
| 174 | if bf16: |
| 175 | model = Model_DIN_V2_Gru_Vec_attGru_Neg_bf16( |
| 176 | n_uid, n_mid, n_cat, EMBEDDING_DIM, HIDDEN_SIZE, |
| 177 | ATTENTION_SIZE) |
| 178 | else: |
| 179 | model = Model_DIN_V2_Gru_Vec_attGru_Neg(n_uid, n_mid, n_cat, |
| 180 | EMBEDDING_DIM, HIDDEN_SIZE, |
| 181 | ATTENTION_SIZE) |
| 182 | |
| 183 | if ev: |
| 184 | sess.run(ops.get_collection(ops.GraphKeys.EV_INIT_VAR_OPS)) |
no test coverage detected