Use pseudo coordinates with gradients about pytorch_spline_conv HOT 4 CLOSED

import torch
import torch.nn.functional as F
from torch_geometric.nn import SplineConv
import torch_geometric
from torch_geometric.data import DataLoader, Data

n_epochs = 2
train_dataset = []
for i in range(10):
    x = torch.ones((3,1))
    edge_index = torch.tensor([[0, 1, 1, 2],
                            [1, 0, 2, 1]], dtype=torch.long)
    pos = torch.rand((3,3))
    edge_attr = torch.rand((4,3))
    target = torch.randint(0,2, (3,34))
    data = Data(x = x, edge_index = edge_index, edge_attr = edge_attr, pos = pos, target = target)
    train_dataset.append(data)

class Net(torch.nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.stn_conv1 = SplineConv(1, 32, dim=3, kernel_size=5, aggr='add')
        self.stn_conv2 = SplineConv(32, 64, dim=3, kernel_size=5, aggr='add')
        self.stn_lin1 = torch.nn.Linear(64, 256)
        self.stn_lin2 = torch.nn.Linear(256, 9)
        self.conv1 = SplineConv(1, 32, dim=3, kernel_size=5, aggr='add')
        self.conv2 = SplineConv(32, 64, dim=3, kernel_size=5, aggr='add')
        self.conv3 = SplineConv(64, 64, dim=3, kernel_size=5, aggr='add')
        self.lin1 = torch.nn.Linear(64, 256)
        self.lin2 = torch.nn.Linear(256, 34) 

    def forward(self, data):
        x, edge_index, pseudo= data.clone().x, data.clone().edge_index, data.clone().edge_attr
        x = F.elu(self.stn_conv1(x, edge_index, pseudo))
        x = F.elu(self.stn_conv2(x, edge_index, pseudo))
        x = F.elu(self.stn_lin1(x))
        x = F.dropout(x, training=self.training)
        x = self.stn_lin2(x)
        x = torch.mean(x, dim=0)
        x = F.normalize(x.view(9), p=2, dim=0)
        rotation_mat = x.view(3,3)
        update_pos = torch.matmul(data.pos,rotation_mat).squeeze(0) 
        update_x = data.x.new_ones(data.x.size())
        row, col = data.edge_index 
        cart = update_pos[col] - update_pos[row] 
        #cart = F.softmax(cart, dim=0).detach() # This passes through but the gradients don't flow back to stn_conv1,..., stn_lin2
        cart = F.softmax(cart, dim=0)
        x = F.elu(self.conv1(update_x, edge_index, cart))
        x = F.elu(self.conv2(x, edge_index, cart))
        x = F.elu(self.conv3(x, edge_index, cart))
        x = F.elu(self.lin1(x))
        x = F.dropout(x, training=self.training)
        x = self.lin2(x)
        return (x)

def train(epoch):
    model.train()
    total_loss = 0
    for data in train_loader:
        target = data.target.to(device).float()
        pred = model(data.to(device))
        optimizer.zero_grad()
        loss = criterion(pred, target)
        loss.backward()
        optimizer.step()
        # print("---gradients after backprop---")
        # for n, p in model.named_parameters():
        #     print(n, p.grad)
        total_loss += loss.item()
    return(total_loss/len(train_dataset))

train_loader = DataLoader(train_dataset, batch_size=1, shuffle=True)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = Net().to(device).float()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
criterion = torch.nn.BCEWithLogitsLoss()

for epoch in range(1, n_epochs):
    train_loss = train(epoch)	
    print('Epoch: {:02d}, train_loss: {:.4f}'.format(epoch, train_loss))