1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104
| import torch import torch.nn as nn
class PostureEstimationModel(nn.Module): """ 基于压力矩阵的姿态估计模型 输入: (batch, 1, 32, 32) 压力矩阵 输出: (batch, num_classes) 姿态分类 """ def __init__(self, num_classes=5): super().__init__() self.features = nn.Sequential( nn.Conv2d(1, 32, kernel_size=3, padding=1), nn.BatchNorm2d(32), nn.ReLU(), nn.MaxPool2d(2), nn.Conv2d(32, 64, kernel_size=3, padding=1), nn.BatchNorm2d(64), nn.ReLU(), nn.MaxPool2d(2), nn.Conv2d(64, 128, kernel_size=3, padding=1), nn.BatchNorm2d(128), nn.ReLU(), nn.AdaptiveAvgPool2d((4, 4)) ) self.classifier = nn.Sequential( nn.Flatten(), nn.Linear(128 * 4 * 4, 256), nn.ReLU(), nn.Dropout(0.5), nn.Linear(256, 64), nn.ReLU(), nn.Linear(64, num_classes) ) self.classes = [ 'NORMAL', 'RECLINED', 'LEANING_FORWARD', 'LEANING_SIDEWAYS', 'EMPTY' ] def forward(self, x): x = self.features(x) x = self.classifier(x) return x def predict(self, pressure_matrix): """预测姿态""" self.eval() with torch.no_grad(): if isinstance(pressure_matrix, np.ndarray): x = torch.from_numpy(pressure_matrix).float() else: x = pressure_matrix x = x / (x.max() + 1e-6) if x.dim() == 2: x = x.unsqueeze(0).unsqueeze(0) elif x.dim() == 3: x = x.unsqueeze(0) output = self(x) probs = torch.softmax(output, dim=1) pred_class = torch.argmax(probs, dim=1) return { 'class': self.classes[pred_class.item()], 'confidence': probs[0, pred_class].item(), 'probabilities': {self.classes[i]: probs[0, i].item() for i in range(len(self.classes))} }
def train_model(): """模型训练流程""" model = PostureEstimationModel() criterion = nn.CrossEntropyLoss() optimizer = torch.optim.Adam(model.parameters(), lr=0.001) print("模型结构:", model) print("参数量:", sum(p.numel() for p in model.parameters()))
|