def test_spline():
model = SplineCNN(16, 32, dim=3, num_layers=2, cat=True, lin=True,
dropout=0.5)
assert model.__repr__() == ('SplineCNN(16, 32, dim=3, num_layers=2, '
'cat=True, lin=True, dropout=0.5)')
x = torch.randn(100, 16)
edge_index = torch.randint(100, (2, 400), dtype=torch.long)
edge_attr = torch.rand((400, 3))
for cat, lin in product([False, True], [False, True]):
model = SplineCNN(16, 32, 3, 2, cat, lin, 0.5)
out = model(x, edge_index, edge_attr)
assert out.size() == (100, 16 + 2 * 32 if not lin and cat else 32)
assert out.size() == (100, model.out_channels)
transform = T.Compose([
T.Constant(),
T.KNNGraph(k=8),
T.Cartesian(),
])
train_dataset = RandomGraphDataset(30, 60, 0, 20, transform=transform)
train_loader = DataLoader(train_dataset, args.batch_size, shuffle=True,
follow_batch=['x_s', 'x_t'])
path = osp.join('..', 'data', 'PascalPF')
test_datasets = [PascalPF(path, cat, transform) for cat in PascalPF.categories]
device = 'cuda' if torch.cuda.is_available() else 'cpu'
psi_1 = SplineCNN(1, args.dim, 2, args.num_layers, cat=False, dropout=0.0)
psi_2 = SplineCNN(args.rnd_dim, args.rnd_dim, 2, args.num_layers, cat=True,
dropout=0.0)
model = DGMC_modified_v2(psi_1, psi_2, num_steps=args.num_steps).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
def train():
model.train()
total_loss = total_examples = total_correct = 0
for i, data in enumerate(train_loader):
optimizer.zero_grad()
data = data.to(device)
S_L = model(data.x_s, data.edge_index_s, data.edge_attr_s,
data.x_s_batch, data.x_t, data.edge_index_t,
pre_filter=pre_filter2
if category in ['car', 'motorbike'] else pre_filter1)
pretrain_datasets += [ValidPairDataset(dataset, dataset, sample=True)]
pretrain_dataset = torch.utils.data.ConcatDataset(pretrain_datasets)
pretrain_loader = DataLoader(pretrain_dataset,
args.batch_size,
shuffle=True,
follow_batch=['x_s', 'x_t'])
path = osp.join('..', 'data', 'WILLOW')
datasets = [WILLOW(path, cat, transform) for cat in WILLOW.categories]
device = 'cuda' if torch.cuda.is_available() else 'cpu'
psi_1 = SplineCNN(dataset.num_node_features,
args.dim,
dataset.num_edge_features,
args.num_layers,
cat=False,
dropout=0.5)
psi_2 = SplineCNN(args.rnd_dim,
args.rnd_dim,
dataset.num_edge_features,
args.num_layers,
cat=True,
dropout=0.0)
model = DGMC_modified(psi_1, psi_2, num_steps=args.num_steps).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
def generate_voc_y(y_col):
y_row = torch.arange(y_col.size(0), device=device)
return torch.stack([y_row, y_col], dim=0)