nn.Linear(args.hs_1, args.d)
)
aggr_net = nn.Sequential(
nn.Linear(args.d+1, args.hs_1), nn.Tanh(),
nn.Linear(args.hs_1, args.hs_1), nn.Tanh(),
nn.Linear(args.hs_1, args.hs_1), nn.Tanh(),
nn.Linear(args.hs_1, 1)
)
model = Model(aggr_net, msg_net)
model.apply(utils.weights_init)
F = DynamicsFunction(model).to(device)
print("Num. of params: {:d}".format(utils.get_parameters_count(model)))
data = utils.read_pickle(['t', 'x', 'u'], args.data_path)
dataset = utils.generate_torchgeom_dataset(data, sig=0.0)
if args.batch_size is None:
batch_size = len(dataset)
else:
batch_size = args.batch_size
loader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
optimizer = optim.Rprop(F.parameters(), lr=args.lr, step_sizes=(1e-8, 10.))
loss_fn = nn.MSELoss()
# Training
ts = dataset[0].t.shape[0] # assumes the same time grid for all sim-s.
print(dataset[0].t)
for epoch in range(args.epochs):
msg_net = nn.Sequential(nn.Linear(6, args.hs_1), nn.Tanh(),
nn.Linear(args.hs_1, args.hs_1), nn.Tanh(),
nn.Linear(args.hs_1, args.hs_1), nn.Tanh(),
nn.Linear(args.hs_1, args.d))
aggr_net = nn.Sequential(nn.Linear(args.d + 2, args.hs_1), nn.Tanh(),
nn.Linear(args.hs_1, args.hs_1), nn.Tanh(),
nn.Linear(args.hs_1, args.hs_1), nn.Tanh(),
nn.Linear(args.hs_1, 2))
model = Model(aggr_net, msg_net)
F = DynamicsFunction(model).to(device)
F.load_state_dict(torch.load(args.model_path, map_location=device))
data = utils.read_pickle(['t', 'x', 'u'], args.data_path)
dataset = utils.generate_torchgeom_dataset(data)
loader = DataLoader(dataset, batch_size=1, shuffle=False)
# Loss
loss_fn = nn.MSELoss()
# Testing
diffs_over_time = []
losses = torch.zeros(len(loader))
inds_of_sims_to_show = set([0])
with torch.no_grad():
for i, dp in enumerate(loader):
edge_index = dp.edge_index
pos = dp.pos