使用resnet18预训练模型做CIFAR10

最新推荐文章于 2023-08-31 23:24:14 发布

Kim_Linshuo

最新推荐文章于 2023-08-31 23:24:14 发布

阅读量868

点赞数

文章标签：深度学习 pytorch 人工智能

本文链接：https://blog.csdn.net/Kim_Linshuo/article/details/126930334

版权

import torchvision.models
import torch
import time
from torch import nn
from torchvision.models import resnet18, ResNet18_Weights
from torch.utils.data import DataLoader
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
data_transform=torchvision.transforms.ToTensor()
train_set=torchvision.datasets.CIFAR10(root="./dataset",train=True,download=False,transform=data_transform)
test_set=torchvision.datasets.CIFAR10(root="./dataset",train=False,download=False,transform=data_transform)
train_dataset=DataLoader(train_set,batch_size=64,shuffle=True)
test_dataset=DataLoader(test_set,batch_size=64)
model=torchvision.models.resnet18(pretrained=True)
num_features=model.fc.in_features
model.fc=nn.Linear(num_features,10)
model=model.to(device)
leng=len(train_set)
optimizer1=torch.optim.SGD(model.parameters(),lr=0.01)
loss=nn.CrossEntropyLoss().to(device)
for epoch in range(10):
    losssum=0.0
    total=0
    accuracy=0.0
    for i,(images,labels) in enumerate(train_dataset):
        optimizer1.zero_grad()
        imgs=images.to(device)
        labels=labels.to(device)
        optimizer1.zero_grad()
        output=model.forward(imgs)
        lossnum=loss(output,labels)
        lossnum.backward()
        optimizer1.step()
        losssum=losssum+lossnum
        accuracy+=(output.argmax(1)==labels).sum()
#     print("Epoch: {} 次的准确率为{}".format(epoch,accuracy/len(train_set)))
with torch.no_grad():
    accuracy=0.0
    for data in test_dataset:
        imgs,labels=data
        imgs=imgs.to(device)
        labels=labels.to(device)
        output=model.forward(imgs)
        lossnum=loss(output,labels)
        accuracy+=(output.argmax(1)==labels).sum()
    print("测试集上的准确率为{}".format(accuracy/len(test_set)))