113 lines
4.0 KiB
Python
113 lines
4.0 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
|
|
class CondGraspNet(nn.Module):
|
|
def __init__(self):
|
|
super(CondGraspNet, self).__init__()
|
|
|
|
# === 1. 定义输入维度 ===
|
|
# 触觉特征: 12维 (3指 * 2单元 * 2分量)
|
|
self.tactile_dim = 12
|
|
# 构型特征: 3维 (One-Hot编码: [1,0,0], [0,1,0], [0,0,1])
|
|
self.config_dim = 3
|
|
|
|
input_total_dim = self.tactile_dim + self.config_dim # 15维
|
|
|
|
# === 2. 定义网络层 (MLP结构) ===
|
|
|
|
# Layer 1: 特征融合层
|
|
# 将触觉信息和构型信息混合
|
|
self.fc1 = nn.Linear(input_total_dim, 64)
|
|
self.bn1 = nn.BatchNorm1d(64) # 批归一化: 防止梯度消失，加速训练
|
|
|
|
# Layer 2: 非线性映射层
|
|
# 增加网络宽度，拟合复杂的力学关系
|
|
self.fc2 = nn.Linear(64, 128)
|
|
self.bn2 = nn.BatchNorm1d(128)
|
|
|
|
# Layer 3: 特征压缩层
|
|
self.fc3 = nn.Linear(128, 64)
|
|
|
|
# Layer 4: 输出层 (Regression Head)
|
|
# 输出3个值: [Delta_X, Delta_Y, Delta_Theta]
|
|
self.output = nn.Linear(64, 3)
|
|
|
|
# === 3. 权重初始化 (Xavier) ===
|
|
# 这一步对小数据集训练非常重要，能让模型收敛得更快
|
|
self._init_weights()
|
|
|
|
def _init_weights(self):
|
|
for m in self.modules():
|
|
if isinstance(m, nn.Linear):
|
|
nn.init.xavier_uniform_(m.weight)
|
|
if m.bias is not None:
|
|
nn.init.constant_(m.bias, 0)
|
|
|
|
def forward(self, tactile_data, config_id_idx):
|
|
"""
|
|
前向传播函数
|
|
:param tactile_data: [Batch_Size, 12] 的触觉数据张量
|
|
:param config_id_idx: [Batch_Size] 的构型索引 (例如 [0, 2, 1...])
|
|
:return: [Batch_Size, 3] 的预测偏差
|
|
"""
|
|
|
|
# Step 1: 处理构型 ID (One-Hot Encoding)
|
|
# 必须把整数 ID (0,1,2) 变成向量 ([1,0,0]...) 才能喂给神经网络
|
|
batch_size = tactile_data.size(0)
|
|
|
|
# 创建一个全0的容器
|
|
config_one_hot = torch.zeros(batch_size, self.config_dim).to(tactile_data.device)
|
|
|
|
# 使用 scatter_ 方法进行填充
|
|
# config_id_idx 需要升维: [Batch] -> [Batch, 1]
|
|
config_one_hot.scatter_(1, config_id_idx.unsqueeze(1).long(), 1)
|
|
|
|
# Step 2: 特征拼接 (Concatenate)
|
|
# 将触觉数据和构型向量拼在一起 -> [Batch, 15]
|
|
x = torch.cat((tactile_data, config_one_hot), dim=1)
|
|
|
|
# Step 3: 通过隐藏层
|
|
x = F.relu(self.bn1(self.fc1(x))) # Linear -> BN -> ReLU
|
|
x = F.relu(self.bn2(self.fc2(x))) # Linear -> BN -> ReLU
|
|
x = F.relu(self.fc3(x)) # Linear -> ReLU (最后一层通常不用BN)
|
|
|
|
# Step 4: 输出结果
|
|
prediction = self.output(x)
|
|
|
|
return prediction
|
|
|
|
|
|
# === 单元测试 (Unit Test) ===
|
|
# 运行此文件，检查网络结构和输入输出形状是否正确
|
|
if __name__ == "__main__":
|
|
print("Testing CondGraspNet Model...")
|
|
|
|
# 1. 实例化模型
|
|
model = CondGraspNet()
|
|
print(f"Model Structure:\n{model}")
|
|
|
|
# 2. 创建模拟输入数据 (Batch Size = 8)
|
|
# 模拟8条触觉数据 (随机数)
|
|
fake_tactile = torch.randn(8, 12)
|
|
# 模拟8个构型ID (随机 0, 1, 2)
|
|
fake_config = torch.tensor([0, 0, 1, 1, 2, 2, 0, 2], dtype=torch.long)
|
|
|
|
# 3. 前向推理
|
|
print("\nProcessing forward pass...")
|
|
try:
|
|
output = model(fake_tactile, fake_config)
|
|
|
|
# 4. 验证结果
|
|
print("Input Shape (Tactile):", fake_tactile.shape)
|
|
print("Output Shape (Pred): ", output.shape) # 期望是 [8, 3]
|
|
print("\nSample Prediction (Row 0):")
|
|
print(f"Delta X: {output[0][0].item():.4f} mm")
|
|
print(f"Delta Y: {output[0][1].item():.4f} mm")
|
|
print(f"Delta θ: {output[0][2].item():.4f} deg")
|
|
|
|
if output.shape == (8, 3):
|
|
print("\n✅ 测试通过：网络维度正确！")
|
|
except Exception as e:
|
|
print(f"\n❌ 测试失败：{e}") |