def train():
dataset = trainDataset()
model = MainModel(dataset.vocab_size()[0], dataset.vocab_size()[1])
#model.load_state_dict(torch.load("model2.pth"))
model = model.cuda()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
dataloader = torch.utils.data.DataLoader(dataset,
batch_size=1,
pin_memory=True)
for epoch in range(EPOCHS):
print(f"EPOCH: {epoch + 1}/{EPOCHS}")
losses = []
for idx, data in tqdm(enumerate(dataloader)):
outputs = model(data["source"].cuda(), data["target"].cuda(),
data["alignment"].cuda())
loss = torch.nn.functional.binary_cross_entropy(
outputs.view(-1), data["predictions"].cuda().view(-1).float())
# print(loss.item())
optimizer.zero_grad()
loss.backward()
optimizer.step()
losses.append(loss.detach())
print(f"Mean Loss for Epoch: {epoch} is {sum(losses) / len(losses)}")
torch.save(model.state_dict(), f"model_lstm.pth")
batch_id += 1
print("g_loss: " + str(g_avg_loss))
logging.info("g_loss: " + str(g_avg_loss))
# Testing
network.cpu().eval()
p_1, r_1, f1_1 = evaluator.sparse_network_reconstruct(network, args.pos_size[0])
p_2, r_2, f1_2 = evaluator.sparse_network_reconstruct(network, args.pos_size[1])
p_3, r_3, f1_3 = evaluator.sparse_network_reconstruct(network, args.pos_size[2])
print(p_3, r_3, f1_3)
if f1_3 > max_f1:
max_f1 = f1_3
max_p = p_3
max_r = r_3
torch.save(network.state_dict(), "model/" + model_dir + ".pt")
network.cuda(args.cuda)
network.generator.sample_linear.cpu()
print("Network Reconstruction Results (K=" + str(args.pos_size[0]) + "): ")
print("Prec: " + str(p_1) + " Rec: " + str(r_1) + " F1: " + str(f1_1))
print("Network Reconstruction Results (K=" + str(args.pos_size[1]) + "): ")
print("Prec: " + str(p_2) + " Rec: " + str(r_2) + " F1: " + str(f1_2))
print("Network Reconstruction Results (K=" + str(args.pos_size[2]) + "): ")
print("Prec: " + str(max_p) + " Rec: " + str(max_r) + " F1: " + str(max_f1))
logging.info("K=" + str(args.pos_size[0]) + "Prec: " + str(p_1) + " Rec: " + str(r_1) + " F1: " + str(f1_1))
logging.info("K=" + str(args.pos_size[1]) + "Prec: " + str(p_2) + " Rec: " + str(r_2) + " F1: " + str(f1_2))
logging.info("K=" + str(args.pos_size[2]) + "Prec: " + str(max_p) + " Rec: " + str(max_r) + " F1: " + str(max_f1))
accu, mle = evaluator.sparse_seq_predict(network, test_marker, test_time, test_mask, adj_list, 1)
print("accu", accu)
print("mle", mle)
logging.info("accu: " + str(accu))
def main(train_dir,
val_dir,
checkpoint_dir,
batch_size,
image_size=512,
num_epochs=10,
checkpoint_name=None,
num_workers=1,
pin_memory=True,
log_dir="logs",
model_name=None,
train_csv=None,
val_csv=None):
# declare datasets
train_ds = DataFolder(root_dir=train_dir,
transform=transform(image_size, is_training=True),
csv_path=train_csv)
val_ds = DataFolder(root_dir=val_dir,
transform=transform(image_size, is_training=False),
csv_path=val_csv)
train_loader = DataLoader(train_ds,
batch_size=batch_size,
num_workers=num_workers,
pin_memory=pin_memory,
shuffle=True)
val_loader = DataLoader(val_ds,
batch_size=batch_size,
num_workers=num_workers,
pin_memory=pin_memory,
shuffle=True)
#init model
model = MainModel(128, model_name)
# configure parameter
loss_fn = nn.CrossEntropyLoss()
model = model.to(device)
optimizer = optim.Adam(model.parameters(), lr=1e-4)
scaler = torch.cuda.amp.GradScaler()
# checkpoint = {'state_dict': model.state_dict(), 'optimizer': optimizer.state_dict()}
# save_checkpoint(checkpoint, os.path.join(checkpoint_dir, f"checkpoint_initialilze.pth.tar"))
# return
if checkpoint_name:
ckp_path = os.path.join(checkpoint_dir, checkpoint_name)
load_checkpoint(torch.load(ckp_path), model, optimizer)
check_accuracy(val_loader, model, device)
#training
for epoch in range(num_epochs):
train_fn(train_loader,
model,
optimizer,
loss_fn,
scaler,
device,
epoch,
log_dir=log_dir)
check_accuracy(val_loader, model, device)
checkpoint = {
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict()
}
save_checkpoint(
checkpoint,
os.path.join(checkpoint_dir, f"checkpoint_{epoch}.pth.tar"))