<filename>train.py
from tgn_modules import TGN, Data, DataWrapper
import torch
import poptorch
from tqdm import tqdm


data_path = 'data/JODIE'
memory_dim = time_dim = embedding_dim = 100
raw_msg_dim = 172
num_nodes = 9227


opts = poptorch.Options()
train_data = poptorch.DataLoader(
    options=opts,
    dataset=DataWrapper(Data(data_path, torch.float32), 'train'),
    batch_size=1,
)

tgn = TGN(
    num_nodes,
    raw_msg_dim,
    memory_dim,
    time_dim,
    embedding_dim,
    dtype=torch.float32,
)

tgn.train()
optim = poptorch.optim.Adam(tgn.parameters(), lr=1e-4)
tgn = poptorch.trainingModel(tgn, optimizer=optim)

for epoch in range(1, 51):
    print(f"Training epoch: {epoch}")
    
    for batch in tqdm(train_data):
        batch = {nm:x.squeeze(0) for nm, x in batch.items()}
        loss = tgn(**batch)
        tgn.memory.detach()
    
    print(f"Loss for epoch {epoch}: {float(loss):.3f}")



<filename>train.py
from tgn_modules import TGN, Data, DataWrapper
import torch
import poptorch
from tqdm import tqdm


data_path = 'data/JODIE'
memory_dim = time_dim = embedding_dim = 100
raw_msg_dim = 172
num_nodes = 9227


opts = poptorch.Options()
train_data = poptorch.DataLoader(
    options=opts,
    dataset=DataWrapper(Data(data_path, torch.float32), 'train'),
    batch_size=1,
)

tgn = TGN(
    num_nodes,
    raw_msg_dim,
    memory_dim,
    time_dim,
    embedding_dim,
    dtype=torch.float32,
)

tgn.train()
optim = poptorch.optim.Adam(tgn.parameters(), lr=1e-4)
tgn = poptorch.trainingModel(tgn, optimizer=optim)

for epoch in range(1, 51):
    print(f"Training epoch: {epoch}")
    
    for batch in tqdm(train_data):
        batch = {nm:x.squeeze(0) for nm, x in batch.items()}
        loss = tgn(**batch)
        tgn.memory.detach()
    
    print(f"Loss for epoch {epoch}: {float(loss):.3f}")


