pytorch2.x 官方quickstart测试

1.本地环境

D:\python2023>nvidia-smi
Thu Jul 27 23:27:45 2023
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 497.29       Driver Version: 497.29       CUDA Version: 11.5     |
|-------------------------------+----------------------+----------------------+
| GPU  Name            TCC/WDDM | Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  NVIDIA GeForce ... WDDM  | 00000000:03:00.0  On |                  N/A |
| 27%   36C    P8     8W / 120W |    397MiB /  3072MiB |      1%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+

+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|    0   N/A  N/A      1288    C+G   Insufficient Permissions        N/A      |
|    0   N/A  N/A      3444    C+G   ...y\ShellExperienceHost.exe    N/A      |
|    0   N/A  N/A      7420    C+G   ...nputApp\TextInputHost.exe    N/A      |
|    0   N/A  N/A      7896    C+G   C:\Windows\explorer.exe         N/A      |
|    0   N/A  N/A      8392    C+G   ...b3d8bbwe\WinStore.App.exe    N/A      |
|    0   N/A  N/A      8872    C+G   ...5n1h2txyewy\SearchApp.exe    N/A      |
|    0   N/A  N/A     10860    C+G   ...lPanel\SystemSettings.exe    N/A      |
|    0   N/A  N/A     11536    C+G   ...se6\Application\360se.exe    N/A      |
|    0   N/A  N/A     14264    C+G   ...\qbblinktrial\browser.exe    N/A      |
+-----------------------------------------------------------------------------+

D:\python2023>gcc --version
gcc (x86_64-posix-sjlj-rev0, Built by MinGW-W64 project) 8.1.0
Copyright (C) 2018 Free Software Foundation, Inc.
This is free software; see the source for copying conditions.  There is NO
warranty; not even for MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.


D:\python2023>python --version
Python 3.8.5

D:\python2023>

2.安装pytorch (Windows GPU版本）

！注意 安装时一定要指定–index-url https://download.pytorch.org/whl/torch/ ，否则安装的是cpu版本，可以访问https://download.pytorch.org/whl/torch/，找到需要的版本如torch-2.0.1+cu117-cp38-cp38-win_amd64.whl 用迅雷下载比较快

pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu117 

3. 官方quickstart

按官方 quickstart拼起來的代码,如果有GPU且安装的GPU版本pytorch则跑GPU上否则CPU（所有CPU），本地测试20CPU与1 个geforce GPU耗时差不多20s

import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision.transforms import ToTensor

#####################
# Download training data from open datasets.
training_data = datasets.FashionMNIST(
    root="data",
    train=True,
    download=True,
    transform=ToTensor(),
)

# Download test data from open datasets.
test_data = datasets.FashionMNIST(
    root="data",
    train=False,
    download=True,
    transform=ToTensor(),
)
#####################
batch_size = 64

# Create data loaders.
train_dataloader = DataLoader(training_data, batch_size=batch_size)
test_dataloader = DataLoader(test_data, batch_size=batch_size)
for X, y in test_dataloader:
    print(f"Shape of X [N, C, H, W]: {X.shape}")
    print(f"Shape of y: {y.shape} {y.dtype}")
    break
#####################
# Get cpu, gpu or mps device for training.
device = (
    "cuda"
    if torch.cuda.is_available()
    else "mps"
    if torch.backends.mps.is_available()
    else "cpu"
)
print(f"Using {device} device")


# Define model
class NeuralNetwork(nn.Module):
    def __init__(self):
        super().__init__()
        self.flatten = nn.Flatten()
        self.linear_relu_stack = nn.Sequential(
            nn.Linear(28 * 28, 512),
            nn.ReLU(),
            nn.Linear(512, 512),
            nn.ReLU(),
            nn.Linear(512, 10)
        )

    def forward(self, x):
        x = self.flatten(x)
        logits = self.linear_relu_stack(x)
        return logits


model = NeuralNetwork().to(device)
print(model)
#######################################
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)


########################################
def train(dataloader, model, loss_fn, optimizer):
    size = len(dataloader.dataset)
    model.train()
    for batch, (X, y) in enumerate(dataloader):
        X, y = X.to(device), y.to(device)

        # Compute prediction error
        pred = model(X)
        loss = loss_fn(pred, y)

        # Backpropagation
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

        if batch % 100 == 0:
            loss, current = loss.item(), (batch + 1) * len(X)
            print(f"loss: {loss:>7f}  [{current:>5d}/{size:>5d}]")


#######################################
def test(dataloader, model, loss_fn):
    size = len(dataloader.dataset)
    num_batches = len(dataloader)
    model.eval()
    test_loss, correct = 0, 0
    with torch.no_grad():
        for X, y in dataloader:
            X, y = X.to(device), y.to(device)
            pred = model(X)
            test_loss += loss_fn(pred, y).item()
            correct += (pred.argmax(1) == y).type(torch.float).sum().item()
    test_loss /= num_batches
    correct /= size
    print(f"Test Error: \n Accuracy: {(100 * correct):>0.1f}%, Avg loss: {test_loss:>8f} \n")


if __name__ == '__main__':
    epochs = 5
    for t in range(epochs):
    	start = time.time()
        print(f"Epoch {t + 1}\n-------------------------------")
        train(train_dataloader, model, loss_fn, optimizer)
        test(test_dataloader, model, loss_fn)
        end = time.time()
        print(f"epoch Done:{end-start}")
    print("Done!")

D:\python2023>nvidia-smi
Fri Jul 28 00:42:59 2023
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 497.29       Driver Version: 497.29       CUDA Version: 11.5     |
|-------------------------------+----------------------+----------------------+
| GPU  Name            TCC/WDDM | Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  NVIDIA GeForce ... WDDM  | 00000000:03:00.0  On |                  N/A |
| 27%   37C    P5     9W / 120W |   1007MiB /  3072MiB |      9%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+

+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|    0   N/A  N/A     11536    C+G   ...se6\Application\360se.exe    N/A      |
|    0   N/A  N/A     14264    C+G   ...\qbblinktrial\browser.exe    N/A      |
|    0   N/A  N/A     15648      C   ...ython\Python38\python.exe    N/A      |
+-----------------------------------------------------------------------------+