Upsampling node features

[hkim716]

I would like to upsample the node features by using clustered edge_index from pooling layers.
Current code is to encode my dataset to obtain 2-dimensional latent space, but I would like to add a decoder to make it back to the original graphs.

train_loader = DataLoader(train_dataset, batch_size=128, shuffle=True) 
test_loader = DataLoader(test_dataset, batch_size=128)
d = train_dataset


def normalized_cut_2d(edge_index, pos):
    row, col = edge_index
    edge_attr = torch.norm(pos[row] - pos[col], p=2, dim=1)
    return normalized_cut(edge_index, edge_attr, num_nodes=pos.size(0))


class Net(torch.nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = SplineConv(1, 32, dim=2, kernel_size=5)
        self.conv2 = SplineConv(32, 64, dim=2, kernel_size=5)
        self.fc1 = torch.nn.Linear(64, 128)
        self.fc2 = torch.nn.Linear(128, 2)

    def forward(self, data):

        data.x = data.x.to(torch.float32)
        data.y = data.y.to(torch.float32)
        data.edge_attr = data.edge_attr.to(torch.float32)
        data.pos = data.pos.to(torch.float32)

        data.x = F.tanh(self.conv1(data.x, data.edge_index, data.edge_attr))
        weight = normalized_cut_2d(data.edge_index, data.pos)
        cluster = graclus(data.edge_index, weight, data.x.size(0))

        data.edge_attr = None
        data = max_pool(cluster, data, transform=transform)

        data.x = F.tanh(self.conv2(data.x, data.edge_index, data.edge_attr))
        weight = normalized_cut_2d(data.edge_index, data.pos)
        cluster = graclus(data.edge_index, weight, data.x.size(0))

        x, batch = max_pool_x(cluster, data.x, data.batch)

        x = global_mean_pool(x, batch)
        x = F.tanh(self.fc1(x))
        x = self.fc2(x)
        x = x.view(-1)
        return x.to(torch.float32)

I saw your previous posting about upsampling node features here, #2139, but I'm not sure how to apply your answer code to my case study.
How could I make and fill the blanks with zero vectors to make it back to the original graph?

[rusty1s]

You can upsample using cluster, e.g.:

x = ...
x_unpooled = x[cluster]

[hkim716]

Ok, I tried to make my model simple by using 1 pooling and 1 up-pooling layers. My code is here below.

train_loader = DataLoader(train_dataset, batch_size=10, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=10)
d = train_dataset

class Net(torch.nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = SplineConv(1, 32, dim=2, kernel_size=5)
        self.conv2 = SplineConv(32, 1, dim=2, kernel_size=5)


    def forward(self, data):
        
        data.x = data.x.to(torch.float32)
        data.y = data.y.to(torch.float32)
        data.edge_attr = data.edge_attr.to(torch.float32)
        data.pos = data.pos.to(torch.float32)
        
        data.x = torch.tanh(self.conv1(data.x, data.edge_index, data.edge_attr))

        weight = normalized_cut_2d(data.edge_index, data.pos)
     
        cluster = graclus(data.edge_index, weight, data.x.size(0))
        
        # Down-Pooling
        coarsened_data = max_pool(cluster, data, transform=transform)

        cluster, perm = self.consecutive_cluster(cluster)

        # Up-Pooling
        upsampled_x = coarsened_data.x[cluster]
        x = upsampled_x

        x = self.conv2(x, data.edge_index, data.edge_attr)

        x = x.view(-1)
        
        return x.to(torch.float32)
        
    def consecutive_cluster(self, src):
        unique, inv = torch.unique(src, sorted=True, return_inverse=True)
        perm = torch.arange(inv.size(0), dtype=inv.dtype, device=inv.device)
        perm = inv.new_empty(unique.size(0)).scatter_(0, inv, perm)
        
        return inv, perm

def train(epoch):
    model.train()
    for data in train_loader:
        data = data.to(device)
        optimizer.zero_grad()
        pred = model(data)
        loss = F.mse_loss(pred, data.x.to(torch.float32).view(-1))
        loss.backward()
        optimizer.step()

Question 1:
Here is my code based on your answer, but I'm not sure if I did right on unpooling part to deal with batch for unpooled data.

Question 2:
There is was an error.

RuntimeError: The size of tensor a (740) must match the size of tensor b (23680) at non-singleton dimension 0

x.shape and model(data).shape from forward function is batch size x num_nodes, but data.x.to(torch.float32).shape from print(data.x.to(torch.float32).shape) was [23680], which is num_channels x batch size x num_nodes.

def train(epoch):
    model.train()
    for data in train_loader:
        data = data.to(device)
        optimizer.zero_grad()
        pred = model(data)
        print(data.x.to(torch.float32).shape)
        loss = F.mse_loss(pred, data.x.to(torch.float32).view(-1))
        loss.backward()
        optimizer.step()

I'm not sure how num_channels could be multiplied to data.x.

Long story short, what I'm trying to do is to make it back to the original graph like this :
original num_nodes -> conv1 ->graph clustering ->~ half num_nodes by pooling -> recover num_nodes by unpooling -> conv2 (did not consider batches and num_channels here)

[rusty1s]

I think your upsampling part is correct. Keep in mind that you modify data.x inside the model execution, so this should fix your issue:

recon = data.x
pred = model(data)
loss = F.mse_loss(pred, recon.view(-1))

[hkim716]

Thanks Matt, It worked!
Now I wanted to see trainable parameters based on this model.

total_param = 0
for name, param in model.named_parameters():
    if param.requires_grad:
        print(name, param.data.shape)
        num_param = 1
        for i in range(len(param.data.shape)):
            num_param *= param.data.shape[i]
        total_param += num_param
        print(total_param)

conv1.weight torch.Size([25, 1, 32])
800
conv1.root torch.Size([1, 32])
832
conv1.bias torch.Size([32])
864
conv2.weight torch.Size([25, 32, 1])
1664
conv2.root torch.Size([32, 1])
1696
conv2.bias torch.Size([1])
1697

I was able to see total number of parameters, which was 1697.
However, I'm not sure what 25 means in the conv1.weight torch.Size([25, 1, 32]).
I used kernel size = 5 and I guess 25 came from kernel size*kernel size.
I think SplineConv uses continuous kernels for unstructured data, so I don't fully understand the meaning of the size of kernels.
Is it grid-shape kernel as we use convolutional kernels with 2d image data? or how do they look like?

[rusty1s]

Yes, 25 refers to a 2D kernel size of 5 * 5. Although SplineConv uses a continuous kernel formulation, its kernel is parametrized by a fixed number of weights in each kernel dimension. Figure 6 in the SplineCNN paper shows how such kernels will look like.