import torch import torch.nn as nn import torch.nn.functional as F import os import copy import math import numpy as np from torch_geometric.data import Data, DataLoader from torch_geometric.nn import ( GCNConv, GATConv, global_mean_pool, global_max_pool, JumpingKnowledge, ) from torch.nn import Sequential, Linear, ReLU, BatchNorm1d, Dropout from torch_geometric.nn import BatchNorm def knn(x, k): """ Computes the k-nearest neighbors for each point in x. Args: x (torch.Tensor): The input tensor of shape (batch_size, num_dims, num_points). k (int): The number of nearest neighbors to find. Returns: torch.Tensor: Indices of the k-nearest neighbors for each point, shape (batch_size, num_points, k). """ # Calculate pairwise distance, shape (batch_size, num_points, num_points) inner = -2 * torch.matmul(x.transpose(2, 1), x) xx = torch.sum(x**2, dim=1, keepdim=True) pairwise_distance = -xx - inner - xx.transpose(2, 1) # Retrieve the indices of the k nearest neighbors idx = pairwise_distance.topk(k=k, dim=-1)[1] return idx def get_graph_feature(x, k=20, idx=None): """ Constructs local graph features for each point by finding its k-nearest neighbors and concatenating the relative position vectors. Args: x (torch.Tensor): The input tensor of shape (batch_size, num_dims, num_points). k (int): The number of neighbors to consider for graph construction. idx (torch.Tensor, optional): Precomputed k-nearest neighbor indices. Returns: torch.Tensor: The constructed graph features of shape (batch_size, 2*num_dims, num_points, k). """ batch_size = x.size(0) num_points = x.size(2) x = x.view(batch_size, -1, num_points) # Compute k-nearest neighbors if not provided if idx is None: idx = knn(x, k=k) # Prepare indices for gathering device = x.device idx_base = torch.arange(0, batch_size, device=device).view(-1, 1, 1) * num_points idx = idx + idx_base idx = idx.view(-1) _, num_dims, _ = x.size() x = x.transpose(2, 1).contiguous() # Gather neighbors for each point to construct local regions feature = x.view(batch_size * num_points, -1)[idx, :] feature = feature.view(batch_size, num_points, k, num_dims) # Expand x to match the dimensions for broadcasting subtraction x = x.view(batch_size, num_points, 1, num_dims).repeat(1, 1, k, 1) # Concatenate the original point features with the relative positions to form the graph features feature = torch.cat((feature - x, x), dim=3).permute(0, 3, 1, 2).contiguous() return feature # class Model(nn.Module): # """ # Deep Graph Convolutional Neural Network for Regression Tasks (RegDGCNN) for processing 3D point cloud data. # This network architecture extracts hierarchical features from point clouds using graph-based convolutions, # enabling effective learning of spatial structures. # """ # def __init__(self, args, device): # """ # Initializes the RegDGCNN model with specified configurations. # Args: # args (dict): Configuration parameters including 'k' for the number of neighbors, 'emb_dims' for embedding # dimensions, and 'dropout' rate. # output_channels (int): Number of output channels (e.g., for drag prediction, this is 1). # """ # super(Model, self).__init__() # self.__name__ = 'RegDGCNN' # self.args = args # self.k = 20 # Number of nearest neighbors # # Batch normalization layers to stabilize and accelerate training # self.bn1 = nn.BatchNorm2d(args.n_hidden) # self.bn2 = nn.BatchNorm2d(args.n_hidden * 2) # self.bn3 = nn.BatchNorm2d(args.n_hidden * 2) # self.bn4 = nn.BatchNorm2d(args.n_hidden * 4) # self.bn5 = nn.BatchNorm1d(args.emb_dims) # # EdgeConv layers: Convolutional layers leveraging local neighborhood information # self.conv1 = nn.Sequential(nn.Conv2d(2*(args.fun_dim + args.space_dim), args.n_hidden, kernel_size=1, bias=False), # self.bn1, # nn.LeakyReLU(negative_slope=0.2)) # self.conv2 = nn.Sequential(nn.Conv2d(args.n_hidden * 2, args.n_hidden * 2, kernel_size=1, bias=False), # self.bn2, # nn.LeakyReLU(negative_slope=0.2)) # self.conv3 = nn.Sequential(nn.Conv2d(args.n_hidden * 4, args.n_hidden * 2, kernel_size=1, bias=False), # self.bn3, # nn.LeakyReLU(negative_slope=0.2)) # self.conv4 = nn.Sequential(nn.Conv2d(args.n_hidden * 4, args.n_hidden * 4, kernel_size=1, bias=False), # self.bn4, # nn.LeakyReLU(negative_slope=0.2)) # self.conv5 = nn.Sequential(nn.Conv1d(args.n_hidden * 9, args.emb_dims, kernel_size=1, bias=False), # self.bn5, # nn.LeakyReLU(negative_slope=0.2)) # # Fully connected layers to interpret the extracted features and make predictions # self.linear1 = nn.Linear(args.emb_dims*2, 128, bias=False) # self.bn6 = nn.LayerNorm(128) # self.dp1 = nn.Dropout(p=args.dropout) # self.linear2 = nn.Linear(128, 64) # self.bn7 = nn.LayerNorm(64) # self.dp2 = nn.Dropout(p=args.dropout) # self.linear3 = nn.Linear(64, 32) # self.bn8 = nn.LayerNorm(32) # self.dp3 = nn.Dropout(p=args.dropout) # self.linear4 = nn.Linear(32, 16) # self.bn9 = nn.LayerNorm(16) # self.dp4 = nn.Dropout(p=args.dropout) # self.linear5 = nn.Linear(16, args.out_dim) # The final output layer # def _make_norm_layer(self, channels): # """创建归一化层,支持batch size=1""" # return nn.GroupNorm(num_groups=min(32, channels), num_channels=channels) # def forward(self, x, fx, T=None, geo=None): # """ # Forward pass of the model to process input data and predict outputs. # Args: # x (torch.Tensor): Input tensor representing a batch of point clouds. # Returns: # torch.Tensor: Model predictions for the input batch. # """ # batch_size = x.size(0) # # 兼容 batch_size = 1 输入：去除 batch 维度 # x = torch.cat([x, fx], dim=-1) # x = x.permute(0, 2, 1) # print(f"x:{x.shape}") # # Extract graph features and apply EdgeConv blocks # x = get_graph_feature(x, k=self.k) # (batch_size, 3, num_points) -> (batch_size, 3*2, num_points, k) # print(f"x:{x.shape}") # x = self.conv1(x) # (batch_size, 3*2, num_points, k) -> (batch_size, 256, num_points, k) # # Global max pooling # x1 = x.max(dim=-1, keepdim=False)[0] # (batch_size, 64, num_points, k) -> (batch_size, 64, num_points) # # Repeat the process for subsequent EdgeConv blocks # x = get_graph_feature(x1, k=self.k) # (batch_size, 256, num_points) -> (batch_size, 256*2, num_points, k) # x = self.conv2(x) # (batch_size, 256*2, num_points, k) -> (batch_size, 512, num_points, k) # x2 = x.max(dim=-1, keepdim=False)[0] # (batch_size, 512, num_points, k) -> (batch_size, 512, num_points) # x = get_graph_feature(x2, k=self.k) # (batch_size, 512, num_points) -> (batch_size, 512*2, num_points, k) # x = self.conv3(x) # (batch_size, 512*2, num_points, k) -> (batch_size, 512, num_points, k) # x3 = x.max(dim=-1, keepdim=False)[0] # (batch_size, 512, num_points, k) -> (batch_size, 512, num_points) # x = get_graph_feature(x3, k=self.k) # (batch_size, 512, num_points) -> (batch_size, 512*2, num_points, k) # x = self.conv4(x) # (batch_size, 512*2, num_points, k) -> (batch_size, 1024, num_points, k) # x4 = x.max(dim=-1, keepdim=False)[0] # (batch_size, 1024, num_points, k) -> (batch_size, 1024, num_points) # # Concatenate features from all EdgeConv blocks # x = torch.cat((x1, x2, x3, x4), dim=1) # (batch_size, 256+512+512+1024, num_points) # # Apply the final convolutional block # x = self.conv5(x) # (batch_size, 256+512+512+1024, num_points) -> (batch_size, emb_dims, num_points) # # Combine global max and average pooling features # # (batch_size, emb_dims, num_points) -> (batch_size, emb_dims) # x1 = F.adaptive_max_pool1d(x, 1).view(batch_size, -1) # # (batch_size, emb_dims, num_points) -> (batch_size, emb_dims) # x2 = F.adaptive_avg_pool1d(x, 1).view(batch_size, -1) # x = torch.cat((x1, x2), 1) # (batch_size, emb_dims*2) # print(f"x:{x.shape}") # # Process features through fully connected layers with dropout and batch normalization # x = F.leaky_relu(self.bn6(self.linear1(x)), negative_slope=0.2) # (batch_size, emb_dims*2) -> (batch_size, 128) # x = self.dp1(x) # x = F.leaky_relu(self.bn7(self.linear2(x)), negative_slope=0.2) # (batch_size, 128) -> (batch_size, 64) # x = self.dp2(x) # x = F.leaky_relu(self.bn8(self.linear3(x)), negative_slope=0.2) # (batch_size, 64) -> (batch_size, 32) # x = self.dp3(x) # x = F.leaky_relu(self.bn9(self.linear4(x)), negative_slope=0.2) # (batch_size, 32) -> (batch_size, 16) # x = self.dp4(x) # # Final linear layer to produce the output # x = self.linear5(x) # (batch_size, 16) -> (batch_size, 1) # print(f"x:{x.shape}") # x = x.permute(0, 2, 1) # return x class Model(nn.Module): """ Deep Graph Convolutional Neural Network for Regression Tasks (RegDGCNN) for processing 3D point cloud data. This network architecture extracts hierarchical features from point clouds using graph-based convolutions, enabling effective learning of spatial structures. """ def __init__(self, args, device): """ Initializes the RegDGCNN model with specified configurations. Args: args (dict): Configuration parameters including 'k' for the number of neighbors, 'emb_dims' for embedding dimensions, and 'dropout' rate. output_channels (int): Number of output channels (e.g., for drag prediction, this is 1). """ super(Model, self).__init__() self.__name__ = "RegDGCNN" self.args = args self.k = 40 # Number of nearest neighbors # Batch normalization layers to stabilize and accelerate training self.bn1 = nn.BatchNorm2d(args.n_hidden) self.bn2 = nn.BatchNorm2d(args.n_hidden * 2) self.bn3 = nn.BatchNorm2d(args.n_hidden * 2) self.bn4 = nn.BatchNorm2d(args.n_hidden * 4) self.bn5 = nn.BatchNorm1d(args.emb_dims) # EdgeConv layers: Convolutional layers leveraging local neighborhood information self.conv1 = nn.Sequential( nn.Conv2d( 2 * (args.fun_dim + args.space_dim), args.n_hidden, kernel_size=1, bias=False, ), self.bn1, nn.LeakyReLU(negative_slope=0.2), ) self.conv2 = nn.Sequential( nn.Conv2d(args.n_hidden * 2, args.n_hidden * 2, kernel_size=1, bias=False), self.bn2, nn.LeakyReLU(negative_slope=0.2), ) self.conv3 = nn.Sequential( nn.Conv2d(args.n_hidden * 4, args.n_hidden * 2, kernel_size=1, bias=False), self.bn3, nn.LeakyReLU(negative_slope=0.2), ) self.conv4 = nn.Sequential( nn.Conv2d(args.n_hidden * 4, args.n_hidden * 4, kernel_size=1, bias=False), self.bn4, nn.LeakyReLU(negative_slope=0.2), ) self.conv5 = nn.Sequential( nn.Conv1d(args.n_hidden * 9, args.emb_dims, kernel_size=1, bias=False), self.bn5, nn.LeakyReLU(negative_slope=0.2), ) # Fully connected layers to interpret the extracted features and make predictions self.point_pred = nn.Sequential( nn.Linear(args.emb_dims, 64), nn.ReLU(), nn.Linear(64, args.out_dim) ) def _make_norm_layer(self, channels): return nn.GroupNorm(num_groups=min(32, channels), num_channels=channels) def forward(self, x, fx, T=None, geo=None): """ Forward pass of the model to process input data and predict outputs. Args: x (torch.Tensor): Input tensor representing a batch of point clouds. Returns: torch.Tensor: Model predictions for the input batch. """ batch_size = x.size(0) # 兼容 batch_size = 1 输入：去除 batch 维度 x = torch.cat([x, fx], dim=-1) x = x.permute(0, 2, 1) # Extract graph features and apply EdgeConv blocks x = get_graph_feature( x, k=self.k ) # (batch_size, 3, num_points) -> (batch_size, 3*2, num_points, k) x = self.conv1( x ) # (batch_size, 3*2, num_points, k) -> (batch_size, 256, num_points, k) # Global max pooling x1 = x.max(dim=-1, keepdim=False)[ 0 ] # (batch_size, 64, num_points, k) -> (batch_size, 64, num_points) # Repeat the process for subsequent EdgeConv blocks x = get_graph_feature( x1, k=self.k ) # (batch_size, 256, num_points) -> (batch_size, 256*2, num_points, k) x = self.conv2( x ) # (batch_size, 256*2, num_points, k) -> (batch_size, 512, num_points, k) x2 = x.max(dim=-1, keepdim=False)[ 0 ] # (batch_size, 512, num_points, k) -> (batch_size, 512, num_points) x = get_graph_feature( x2, k=self.k ) # (batch_size, 512, num_points) -> (batch_size, 512*2, num_points, k) x = self.conv3( x ) # (batch_size, 512*2, num_points, k) -> (batch_size, 512, num_points, k) x3 = x.max(dim=-1, keepdim=False)[ 0 ] # (batch_size, 512, num_points, k) -> (batch_size, 512, num_points) x = get_graph_feature( x3, k=self.k ) # (batch_size, 512, num_points) -> (batch_size, 512*2, num_points, k) x = self.conv4( x ) # (batch_size, 512*2, num_points, k) -> (batch_size, 1024, num_points, k) x4 = x.max(dim=-1, keepdim=False)[ 0 ] # (batch_size, 1024, num_points, k) -> (batch_size, 1024, num_points) # Concatenate features from all EdgeConv blocks x = torch.cat( (x1, x2, x3, x4), dim=1 ) # (batch_size, 256+512+512+1024, num_points) # Apply the final convolutional block x = self.conv5( x ) # (batch_size, 256+512+512+1024, num_points) -> (batch_size, emb_dims, num_points) x = x.permute(0, 2, 1) # (batch_size, num_points, emb_dims) out = self.point_pred(x) # 点级预测层，输出 (batch_size, num_points, out_dim) return out