PyTorch在Ubuntu上如何进行模型训练

?Ubuntu???PyTorch??????,??????????:

??PyTorch

1. ?????????: ????CUDA??(????NVIDIA GPU),?????PyTorch????????PyTorch????????????

   # ??,??CUDA 11.7 pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu117 

2. ????: ??Python????????PyTorch??????????

   import torch print(torch.__version__) print(torch.cuda.is_available()) # ???GPU,????True 

?????

1. ?????: ????PyTorch??????,??????????

   from torchvision import datasets, transforms # ?????? transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,)) ]) # ??????? train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform) 

2. ???????: ??DataLoader????????

   from torch.utils.data import DataLoader train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True) 

????

1. ??PyTorch?nn??????:

   import torch.nn as nn import torch.nn.functional as F class Net(nn.Module): def __init__(self): super(Net, self).__init__() self.conv1 = nn.Conv2d(1, 10, kernel_size=5) self.conv2 = nn.Conv2d(10, 20, kernel_size=5) self.conv2_drop = nn.Dropout2d() self.fc1 = nn.Linear(320, 50) self.fc2 = nn.Linear(50, 10) def forward(self, x): x = F.relu(F.max_pool2d(self.conv1(x), 2)) x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2)) x = x.view(-1, 320) x = F.relu(self.fc1(x)) x = F.dropout(x, training=self.training) x = self.fc2(x) return F.log_softmax(x, dim=1) 

????

1. ??????????????:

   model = Net() optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.5) criterion = nn.CrossEntropyLoss() 

2. ??????:

   for epoch in range(10): # ????????? running_loss = 0 for i, data in enumerate(train_loader, 0): # ?????? inputs, labels = data # ???? optimizer.zero_grad() # ???? + ???? + ?? outputs = model(inputs) loss = criterion(outputs, labels) loss.backward() optimizer.step() # ?????? running_loss += loss.item() if i % 100 == 99: # ?100?mini-batches???? print('[%d, ]] loss: %.3f' % (epoch + 1, i + 1, running_loss / 100)) running_loss = 0.0 print('Finished Training') 

???????

1. ????:

   PATH = './cifar_net.pth' torch.save(model.state_dict(), PATH) 

2. ????:

   model = Net() model.load_state_dict(torch.load(PATH)) 

????

• ???????????CUDA?cuDNN,???????GPU?????
• ????????????????????????
• ??????????????,??????????

??????????????,????????????????????