densenet121采用pytorch预训练模型，这里用cifar10作为数据集
import torchvision.models as models import ssl ssl._create_default_https_context = ssl._create_unverified_context import torch import torch.nn as nn import torchvision import torchvision.transforms as transforms import torch.optim as optim device = torch.device("cuda" if torch.cuda.is_available() else "cpu") num_classes = 10 num_epochs = 10 batch_size = 100 learning_rate = 0.001 model = models.densenet121(pretrained=True) for param in model.parameters(): param.requires_grad = True model.fc = nn.Linear(512, 10) model = model.to(device) # Optimize only the classifier transform = transforms.Compose( [transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]) trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform) trainloader = torch.utils.data.DataLoader(trainset, batch_size=100, shuffle=True, num_workers=2) testset = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform) testloader = torch.utils.data.DataLoader(testset, batch_size=100, shuffle=False, num_workers=2) classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck') criterion = nn.CrossEntropyLoss() optimizer = optim.SGD(model.parameters(), lr=1e-2, momentum=0.9) """train the model""" total_step = len(trainloader) for epoch in range(num_epochs): for i,(images,labels) in enumerate(trainloader): images, labels = images.to(device), labels.to(device) outputs = model(images) loss = criterion(outputs,labels) optimizer.zero_grad() loss.backward() optimizer.step() if (i+1) % 100 == 0: print('Epoch [{}/{}], Step[{}/{}],Loss:{:.4f}'\ .format(epoch+1,num_epochs,i+1,total_step,loss.item())) correct = 0 total = 0 with torch.no_grad(): for data in testloader: images, labels = data images, labels = images.to(device), labels.to(device) outputs = model(images) _, predicted = torch.max(outputs.data, 1) total += labels.size(0) correct += (predicted == labels).sum().item() print('Accuracy of the network on the 10000 test images: %d %%' % ( 100 * correct / total)) class_correct = list(0. for i in range(10)) class_total = list(0. for i in range(10)) with torch.no_grad(): for data in testloader: images, labels = data images, labels = images.to(device), labels.to(device) outputs = model(images) _, predicted = torch.max(outputs, 1) c = (predicted == labels).squeeze() for i in range(100): label = labels[i] class_correct[label] += c[i].item() class_total[label] += 1 for i in range(10): print('Accuracy of %5s : - %%' % (\ classes[i], 100 * class_correct[i] / class_total[i]))为了演示，只训练了10轮，效果为55% 各个类的分类精度也给出了。
