randla-net

import torch
import torch.nn as nn
from torch_geometric.nn import MessagePassing
from torch_geometric.utils import add_self_loops, degree

class RandLANetConv(MessagePassing):
    def __init__(self, in_channels, out_channels):
        super(RandLANetConv, self).__init__(aggr='add')
        self.lin = nn.Linear(in_channels, out_channels)

    def forward(self, x, edge_index):
        edge_index, _ = add_self_loops(edge_index, num_nodes=x.size(0))
        row, col = edge_index
        deg = degree(row, x.size(0), dtype=x.dtype)
        deg_inv_sqrt = deg.pow(-0.5)
        norm = deg_inv_sqrt[row] * deg_inv_sqrt[col]
        return self.propagate(edge_index, x=x, norm=norm)

    def message(self, x_j, norm):
        return norm.view(-1, 1) * self.lin(x_j)

class RandLANet(nn.Module):
    def __init__(self, in_channels, hidden_channels, out_channels):
        super(RandLANet, self).__init__()
        self.conv1 = RandLANetConv(in_channels, hidden_channels)
        self.conv2 = RandLANetConv(hidden_channels, out_channels)

    def forward(self, x, edge_index):
        x = self.conv1(x, edge_index)
        x = F.relu(x)
        x = F.dropout(x, p=0.5, training=self.training)
        x = self.conv2(x, edge_index)
        return x

# 示例使用
model = RandLANet(in_channels=3, hidden_channels=64, out_channels=10)