class GCNTrainer(Trainer):
def __init__(self, parameter, emb_matrix=None):
self.parameter = parameter
self.emb_matrix = emb_matrix
self.model = GCN(parameter, emb_matrix=emb_matrix)
self.criterion = nn.CrossEntropyLoss()
self.parameters = [
p for p in self.model.parameters() if p.requires_grad
]
self.optimizer = torch.optim.SGD(self.parameters,
lr=0.5,
weight_decay=0)
def update(self, batch):
inputs, labels = unpack_batch(batch)
self.model.train() #开启dropout
self.optimizer.zero_grad()
logits, pooling_output = self.model(inputs)
loss = self.criterion(logits, labels)
#正则化项
loss += 0.003 * (pooling_output**2).sum(1).mean()
loss_val = loss.item()
# 反向传播
loss.backward()
torch.nn.utils.clip_grad_norm_(self.model.parameters(), 5)
self.optimizer.step()
return loss_val
def predict(self, batch):
inputs, labels = unpack_batch(batch)
self.model.eval() #关闭dropout
orig_idx = batch[-1]
logits, _ = self.model(inputs)
probs = F.softmax(logits, 1).data.cpu().numpy().tolist()
predictions = np.argmax(logits.data.cpu().numpy(), axis=1).tolist()
_, predictions, probs = [list(t) for t in zip(*sorted(zip(orig_idx,\
predictions, probs)))]
return predictions
def train(model_save_path, root_dir, day, train_log_dir, args, layer, dim):
# print params
tab_printer(args)
# load train dataset
dataset = DataSet(os.path.join(root_dir, day)) # TODO: 更改训练数据
# create model on GPU:1
running_context = torch.device("cuda:0")
model = GCN(args, running_context, layer, dim).to(running_context)
optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
model.train()
# train logs
train_log_list = np.zeros((args.epochs, 24), dtype=np.float)
print("\n" + "+" * 5 + " Train on day {} layer {} dim {} ".format(day, layer, dim) + "+" * 5)
best_loss = float("inf")
last_loss_decrease_epoch = 0
stop_flag = False
best_model = None
train_start_time = time.perf_counter()
for epoch in range(args.epochs):
start_time = time.perf_counter()
total_loss = 0.
print("\n" + "+" * 10 + " epoch {:3d} ".format(epoch) + "+" * 10)
for i in range(len(dataset)): # 24 hours
data = dataset[i]
edge_index = data.edge_index.to(running_context)
mask = data.mask.to(running_context)
logits = model(data.inputs, edge_index)
label = data.label.to(running_context, non_blocking=True)
pos_cnt = torch.sum(label == 1)
neg_cnt = torch.sum(label == 0)
weight = torch.tensor([pos_cnt.float(), neg_cnt.float()]) / (neg_cnt + pos_cnt)
loss_fn = nn.CrossEntropyLoss(weight=weight).to(running_context)
loss = loss_fn(logits[mask], label)
total_loss += loss.item()
optimizer.zero_grad()
loss.backward()
optimizer.step()
# acc, tn, fp, fn, tp, white_p, white_r, white_incorrect, black_p, black_r, black_incorrect, macro_f1, micro_f1 \
# = evaluate(logits[mask].max(1)[1], label)
train_log_list[epoch][i] = float(loss.item())
print("hour: {:2d}, loss: {:.4f}".format(i, loss.item()))
if total_loss < best_loss:
best_loss = total_loss
last_loss_decrease_epoch = epoch
best_model = model.state_dict()
else:
not_loss_decrease_epochs = epoch - last_loss_decrease_epoch
if not_loss_decrease_epochs >= args.early_stop:
stop_flag = True
else:
pass
if stop_flag:
print("early stop...")
save_model(best_model, model_save_path)
print("\nepoch: {:3d}, total_loss: {:.4f}, best_loss: {:.4f} time: {:.4f}" \
.format(epoch + 1, total_loss, best_loss, time.perf_counter() - start_time))
print("\ntotal train time: {}".format(time.perf_counter() - train_start_time))
# save model when not early stop
if not stop_flag:
save_model(best_model, model_save_path)
# save train logs to csv file
print("Start to save train log of {}.".format(day))
train_log_cols = ["hour_{}".format(hour) for hour in range(24)]
train_log_df = pd.DataFrame(train_log_list, columns=train_log_cols)
train_log_df.to_csv(train_log_dir, float_format="%.4f", index=None, columns=train_log_cols)
print("Save train log of {} layer {} dim {} successfully.".format(day, layer, dim))
torch.cuda.empty_cache()
model = GCN(hidden_num=1,
hidden_dim=[hidden_dim],
**addi_parameters,
learning_rate=learning_rate,
epochs=epochs,
weight_decay=weight_decay,
early_stopping=early_stopping,
activation_func=activation_func,
dropout_prob=dropout_prob,
bias=bias,
optimizer=optimizer,
name='GCN')
sess = tf.Session()
model.train(sess, undirected, features, y_train, y_val, train_mask,
val_mask, num_featuers_nonzero)
accu, time_used = model.test(sess, undirected, features, y_test, test_mask,
num_featuers_nonzero)
print('test acucracy: ', accu)
#Test GCN finish
#Test MLP
elif model_name == 'mlp':
#Create model
model = MLP(hidden_num=1,
hidden_dim=[hidden_dim],
**addi_parameters,
learning_rate=learning_rate,
epochs=epochs,
