def train_epoch_dense(model, optimizer, device, data_loader, epoch,
batch_size):
model.train()
epoch_loss = 0
epoch_train_mae = 0
nb_data = 0
gpu_mem = 0
optimizer.zero_grad()
for iter, (x_no_edge_feat, x_with_edge_feat,
targets) in enumerate(data_loader):
if x_no_edge_feat is not None:
x_no_edge_feat = x_no_edge_feat.to(device)
if x_with_edge_feat is not None:
x_with_edge_feat = x_with_edge_feat.to(device)
targets = targets.to(device)
scores = model.forward(x_no_edge_feat, x_with_edge_feat)
loss = model.loss(scores, targets)
loss.backward()
if not (iter % batch_size):
torch.nn.utils.clip_grad_norm(model.parameters(), 3)
optimizer.step()
optimizer.zero_grad()
epoch_loss += loss.detach().item()
epoch_train_mae += MAE(scores, targets)
nb_data += targets.size(0)
epoch_loss /= (iter + 1)
epoch_train_mae /= (iter + 1)
return epoch_loss, epoch_train_mae, optimizer
def evaluate_network(model, device, data_loader, epoch):
model.eval()
epoch_test_loss = 0
epoch_test_mae = 0
nb_data = 0
with torch.no_grad():
for iter, (batch_graphs, batch_targets) in enumerate(data_loader):
batch_graphs = batch_graphs.to(device)
batch_x = batch_graphs.ndata['feat'].to(device)
batch_e = batch_graphs.edata['feat'].to(device)
batch_targets = batch_targets.to(device)
try:
batch_lap_pos_enc = batch_graphs.ndata['lap_pos_enc'].to(device)
except:
batch_lap_pos_enc = None
try:
batch_wl_pos_enc = batch_graphs.ndata['wl_pos_enc'].to(device)
except:
batch_wl_pos_enc = None
batch_scores = model.forward(batch_graphs, batch_x, batch_e, batch_lap_pos_enc, batch_wl_pos_enc)
loss = model.loss(batch_scores, batch_targets)
epoch_test_loss += loss.detach().item()
epoch_test_mae += MAE(batch_scores, batch_targets)
nb_data += batch_targets.size(0)
epoch_test_loss /= (iter + 1)
epoch_test_mae /= (iter + 1)
return epoch_test_loss, epoch_test_mae
def train_epoch_sparse(model, optimizer, device, data_loader, epoch):
model.train()
epoch_loss = 0
epoch_train_mae = 0
nb_data = 0
gpu_mem = 0
for iter, (batch_graphs, batch_targets) in enumerate(data_loader):
batch_graphs.to(device)
batch_x = batch_graphs.ndata['feat'].to(device) # num x feat
batch_e = batch_graphs.edata['feat'].to(device)
batch_targets = batch_targets.to(device)
optimizer.zero_grad()
try:
batch_pos_enc = batch_graphs.ndata['pos_enc'].to(device)
sign_flip = torch.rand(batch_pos_enc.size(1)).to(device)
sign_flip[sign_flip >= 0.5] = 1.0
sign_flip[sign_flip < 0.5] = -1.0
batch_pos_enc = batch_pos_enc * sign_flip.unsqueeze(0)
batch_scores = model.forward(batch_graphs, batch_x, batch_e,
batch_pos_enc)
except:
batch_scores = model.forward(batch_graphs, batch_x, batch_e)
loss = model.loss(batch_scores, batch_targets)
loss.backward()
optimizer.step()
epoch_loss += loss.detach().item()
epoch_train_mae += MAE(batch_scores, batch_targets)
nb_data += batch_targets.size(0)
epoch_loss /= (iter + 1)
epoch_train_mae /= (iter + 1)
return epoch_loss, epoch_train_mae, optimizer
def train_epoch_sparse(model, optimizer, device, data_loader, epoch):
model.train()
epoch_loss = 0
epoch_train_mae = 0
nb_data = 0
gpu_mem = 0
for iter, (batch_graphs, batch_targets) in enumerate(data_loader):
batch_x = batch_graphs.ndata['feat'].to(device) # num x feat
batch_e = batch_graphs.edata['feat'].to(device)
batch_targets = batch_targets.to(device)
optimizer.zero_grad()
batch_scores = model.forward(batch_graphs, batch_x, batch_e)
loss = model.loss(batch_scores, batch_targets)
loss.backward()
optimizer.step()
epoch_loss += loss.detach().item()
epoch_train_mae += MAE(batch_scores, batch_targets)
nb_data += batch_targets.size(0)
epoch_loss /= (iter + 1)
epoch_train_mae /= (iter + 1)
return epoch_loss, epoch_train_mae, optimizer
def train_epoch_sparse(model, optimizer, device, data_loader, epoch,
MODEL_NAME):
model.train()
epoch_loss = 0
epoch_train_mae = 0
nb_data = 0
gpu_mem = 0
for iter, (batch_graphs, batch_targets) in enumerate(data_loader):
batch_x = batch_graphs.ndata['feat'].to(device) # num x feat
"""
print('\n\nStart with graph')
#Reformat the node features from 2 to 29 using 1-hot encoding
enc_feat = list()
for feat in batch_x:
np_feat = feat.numpy()
ndata_enc = [0] * 29
ndata_enc[int(np_feat[0])] = 1.0
ndata_enc[28] = np_feat[1]
enc_feat.append(ndata_enc)
batch_x = torch.tensor(enc_feat)
"""
batch_e = batch_graphs.edata['feat'].to(device)
batch_targets = batch_targets.to(device)
optimizer.zero_grad()
try:
batch_pos_enc = batch_graphs.ndata['pos_enc'].to(device)
sign_flip = torch.rand(batch_pos_enc.size(1)).to(device)
sign_flip[sign_flip >= 0.5] = 1.0
sign_flip[sign_flip < 0.5] = -1.0
batch_pos_enc = batch_pos_enc * sign_flip.unsqueeze(0)
batch_scores = model.forward(batch_graphs.to(device), batch_x,
batch_e, batch_pos_enc)
except:
batch_scores = model.forward(batch_graphs.to(device), batch_x,
batch_e)
loss = model.loss(batch_scores, batch_targets)
loss.backward()
optimizer.step()
epoch_loss += loss.detach().item()
epoch_train_mae += MAE(batch_scores, batch_targets)
nb_data += batch_targets.size(0)
epoch_loss /= (iter + 1)
epoch_train_mae /= (iter + 1)
return epoch_loss, epoch_train_mae, optimizer
def evaluate_network_sparse(model, device, data_loader, epoch):
model.eval()
epoch_test_loss = 0
epoch_test_mae = 0
nb_data = 0
with torch.no_grad():
for iter, (batch_graphs, batch_targets) in enumerate(data_loader):
batch_x = batch_graphs.ndata['feat'].float().to(device)
batch_targets = batch_targets.float().to(device)
batch_scores = model.forward(batch_graphs, batch_x, None)
loss = model.loss(batch_scores, batch_targets)
epoch_test_loss += loss.detach().item()
epoch_test_mae += MAE(batch_scores, batch_targets)
nb_data += batch_targets.size(0)
epoch_test_loss /= (iter + 1)
epoch_test_mae /= (iter + 1)
return epoch_test_loss, epoch_test_mae
def evaluate_network_dense(model, device, data_loader, epoch):
model.eval()
epoch_test_loss = 0
epoch_test_mae = 0
nb_data = 0
with torch.no_grad():
for iter, (x_no_edge_feat, x_with_edge_feat, targets) in enumerate(data_loader):
if x_no_edge_feat is not None:
x_no_edge_feat = x_no_edge_feat.to(device)
if x_with_edge_feat is not None:
x_with_edge_feat = x_with_edge_feat.to(device)
targets = targets.to(device)
scores = model.forward(x_no_edge_feat, x_with_edge_feat)
loss = model.loss(scores, targets)
epoch_test_loss += loss.detach().item()
epoch_test_mae += MAE(scores, targets)
nb_data += targets.size(0)
epoch_test_loss /= (iter + 1)
epoch_test_mae /= (iter + 1)
return epoch_test_loss, epoch_test_mae