🧠 PyTorch 入门教程（零基础到训练模型）

pip install torch torchvision torchaudio

conda install pytorch torchvision torchaudio pytorch-cuda=11.8 -c pytorch -c nvidia

import torch

a = torch.tensor([1.0, 2.0, 3.0])
b = torch.ones(3)
c = torch.zeros(3)

print(a + b)
print(a * 2)
print(a.shape)

x = torch.tensor(2.0, requires_grad=True)
y = x ** 2
y.backward()
print(x.grad)  # 输出：4.0

import torch.nn as nn

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc = nn.Linear(10, 1)

    def forward(self, x):
        return self.fc(x)

model = Net()
print(model)

import torch.optim as optim

criterion = nn.MSELoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)

inputs = torch.randn(5, 10)
targets = torch.randn(5, 1)

optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, targets)
loss.backward()
optimizer.step()

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

torch.save(model.state_dict(), 'model.pth')
model.load_state_dict(torch.load('model.pth'))
model.eval()