优化器

  • step方法：利用梯度对参数进行更新
  • 示例代码：
  1. import torch
  2. from torch import nn
  3. import torchvision
  4. from torch.nn import Conv2d, Sequential, MaxPool2d, Flatten, Linear, CrossEntropyLoss
  5. from torch.utils.data import DataLoader
  7. dataset = torchvision.datasets.CIFAR10("./download_data",transform=torchvision.transforms.ToTensor(),
  8.                                        train=False, download=True)
  9. dataloader = DataLoader(dataset,batch_size=1)
  11. class Liucy(nn.Module):
  12.     def __init__(self):
  13.         super(Liucy,self).__init__()
  14.         self.modle = Sequential(
  15.             Conv2d(3, 32, 5, padding=2),
  16.             MaxPool2d(2),
  17.             Conv2d(32, 32, 5, padding=2),
  18.             MaxPool2d(2),
  19.             Conv2d(32, 64, 5, padding=2),
  20.             MaxPool2d(2),
  21.             Flatten(),
  22.             Linear(1024, 64),
  23.             Linear(64, 10)
  24.         )
  26.     def forward(self,x):
  27.         x = self.modle(x)
  28.         return x
  30. liu = Liucy()
  31. cross = CrossEntropyLoss()
  33. # 设置优化器
  34. optim = torch.optim.SGD(liu.parameters(),lr=0.01) #lr为学习速率
  35. for i in range(10):
  36.     lost_value = 0.0
  37.     for data in dataloader:
  38.         img,label = data
  39.         output = liu(img)
  40.         # print(output) #最后会有十个分类
  41.         result = cross(output,label)
  42.         optim.zero_grad()
  43.         result.backward()
  44.         optim.step()
  45.         lost_value = lost_value + result
  46.     print(lost_value)