把神经网络丢到GPU上跑

第一步：

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

GPU