把神经网络丢到GPU上跑

关注

原创最新推荐文章于 2026-02-28 02:54:11 发布 359 阅读 AI 写同款· GEO 优化 ›

Nuvoton端点AI平台解析：Cortex-M55与Ethos-U55让MCU跑神经网络 边缘计算和端侧AI正在重塑嵌入式设备的能力边界，而传统MCU受限于算力与内存，难以在本地执行神经网络推理。为突破这一瓶颈，异构计算架构成为关键：主核负责应用控制，专用NPU专注张量运算，从而在保持低功耗和实时性的同时实现端侧智能。这种技术路径的价值在于，它让传感器、家电、工业设备等最末梢的终端也能离线完成图像分类、语音唤醒和异常检测，无需依赖云端，兼顾隐私与响应速度。Nuvoton推出的Endpoint AI平台，正是基于Cortex-M55与Ethos-U55 NPU的MCU级方案，配合Vela编译器与E 阅读详情

第一步：

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

第二步：

把训练集验证集、net、loss什么的通通

.to(device)

比如

X_train = X_train.to(device)

完整代码如下（因为使用了很多自己的函数，所以仅供参考）：

from torch import nn
import torch

from easier_nn.classic_dataset import fashion_mnist
from easier_nn.train_net import train_net_with_evaluation

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

fm = fashion_mnist()
fm.load_fashion_mnist(flatten=False)
# 将数据移动到 GPU 上
fm.X_train = fm.X_train.to(device)
fm.y_train = fm.y_train.to(device)
fm.X_test = fm.X_test.to(device)
fm.y_test = fm.y_test.to(device)

lr = 0.001  # 学习率
num_epochs = 30  # 迭代周期
batch_size = 64  # 每个小批量样本的数量
train_iter, test_iter = fm.load_dataiter(batch_size=batch_size)
train_iter = [(X.to(device), y.to(device)) for X, y in train_iter]
test_iter = [(X.to(device), y.to(device)) for X, y in test_iter]
net = nn.Sequential(
    nn.Conv2d(1, 6, kernel_size=5, padding=2), nn.Sigmoid(),
    nn.AvgPool2d(kernel_size=2, stride=2),
    nn.Conv2d(6, 16, kernel_size=5), nn.Sigmoid(),
    nn.AvgPool2d(kernel_size=2, stride=2),
    nn.Flatten(),
    nn.Linear(16 * 5 * 5, 120), nn.Sigmoid(),
    nn.Linear(120, 84), nn.Sigmoid(),
    nn.Linear(84, 10))
net.to(device)
loss = nn.CrossEntropyLoss()
loss.to(device)
optimizer = torch.optim.Adam(net.parameters(), lr=lr)

net, _, _, _ = train_net_with_evaluation(fm.X_train, fm.y_train, fm.X_test, fm.y_test, data_iter=train_iter,
                                         test_iter=test_iter, net=net, loss=loss, optimizer=optimizer,
                                         num_epochs=num_epochs, show_interval=2, draw='acc')
fm.predict(net, test_iter, n=18, num_rows=3, num_cols=6)

效果：CPU