手把手教你用Pytorch训练图像分类网络（pytorch实现CNN）

手把手教你用Pytorch训练图像分类网络（pytorch实现CNN）

import torchvisionfrom torch import nnfrom torch.utils.data import DataLoaderimport torchdata_set = torchvision.datasets.CIFAR10(train=False, root='./dataset/', download=True, transform=torchvision.transforms.Compose([ torchvision.transforms.ToTensor(), ]))data_loader = DataLoader(dataset=data_set, batch_size=128)class Model(nn.Module): def __init__(self): super(Model, self).__init__() self.sequential = nn.Sequential( nn.Conv2d(in_channels=3, out_channels=32, kernel_size=5, padding=2), nn.MaxPool2d(kernel_size=2), nn.Conv2d(in_channels=32, out_channels=32, kernel_size=5, padding=2), nn.MaxPool2d(kernel_size=2), nn.Conv2d(in_channels=32, out_channels=64, kernel_size=5, padding=2), nn.MaxPool2d(kernel_size=2), nn.Flatten(), nn.Linear(in_features=1024, out_features=64), nn.Linear(in_features=64, out_features=10) ) def forward(self, input): output = self.sequential(input) return outputmodel = Model()loss = nn.CrossEntropyLoss()optim = torch.optim.SGD(params=model.parameters(), lr=0.01)for epoch in range(20): epoch_loss = 0.0 for data in data_loader: imgs, targets = data output = model(imgs) res_loss = loss(output, targets) # step 1 梯度清0 optim.zero_grad() # step 2 计算梯度 res_loss.backward() # 反向传播 # step 3 梯度更新 optim.step() epoch_loss += res_loss print(epoch_loss)

版权声明：本文内容由网络用户投稿，版权归原作者所有，本站不拥有其著作权，亦不承担相应法律责任。如果您发现本站中有涉嫌抄袭或描述失实的内容，请联系我们jiasou666@gmail.com 处理，核实后本网站将在24小时内删除侵权内容。