测试CUDA
import torch
import numpy as np# check if CUDA is available
train_on_gpu = torch.cuda.is_available()if not train_on_gpu:print('CUDA is not available. Training on CPU ...')
else:print('CUDA is available! Training on GPU ...')加载数据
from torchvision import datasets
import torchvision.transforms as transforms
from torch.utils.data.sampler import SubsetRandomSampler# number of subprocesses to use for data loading
num_workers = 0
# how many samples per batch to load
batch_size = 20
# percentage of training set to use as validation
valid_size = 0.2# convert data to a normalized torch.FloatTensor
transform = transforms.Compose([transforms.RandomHorizontalFlip(), # randomly flip and rotatetransforms.RandomRotation(10),transforms.ToTensor(),transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])# choose the training and test datasets
train_data = datasets.CIFAR10('data', train=True,download=True, transform=transform)
test_data = datasets.CIFAR10('data', train=False,download=True, transform=transform)# obtain training indices that will be used for validation
num_train = len(train_data)
indices = list(range(num_train))
np.random.shuffle(indices)
split = int(np.floor(valid_size * num_train))
train_idx, valid_idx = indices[split:], indices[:split]# define samplers for obtaining training and validation batches
train_sampler = SubsetRandomSampler(train_idx)
valid_sampler = SubsetRandomSampler(valid_idx)# prepare data loaders (combine dataset and sampler)
train_loader = torch.utils.data.DataLoader(train_data, batch_size=batch_size,sampler=train_sampler, num_workers=num_workers)
valid_loader = torch.utils.data.DataLoader(train_data, batch_size=batch_size, sampler=valid_sampler, num_workers=num_workers)
test_loader = torch.utils.data.DataLoader(test_data, batch_size=batch_size, num_workers=num_workers)# specify the image classes
classes = ['airplane', 'automobile', 'bird', 'cat', 'deer','dog', 'frog', 'horse', 'ship', 'truck']可视化训练数据
import matplotlib.pyplot as plt
%matplotlib inline# helper function to un-normalize and display an image
def imshow(img):img = img / 2 + 0.5 # unnormalizeplt.imshow(np.transpose(img, (1, 2, 0))) # convert from Tensor image# obtain one batch of training images
dataiter = iter(train_loader)
images, labels = dataiter.next()
images = images.numpy() # convert images to numpy for display# plot the images in the batch, along with the corresponding labels
fig = plt.figure(figsize=(25, 4))
# display 20 images
for idx in np.arange(20):ax = fig.add_subplot(2, 20/2, idx+1, xticks=[], yticks=[])imshow(images[idx])ax.set_title(classes[labels[idx]])详细查看图像
rgb_img = np.squeeze(images[3])
channels = ['red channel', 'green channel', 'blue channel']fig = plt.figure(figsize = (36, 36))
for idx in np.arange(rgb_img.shape[0]):ax = fig.add_subplot(1, 3, idx + 1)img = rgb_img[idx]ax.imshow(img, cmap='gray')ax.set_title(channels[idx])width, height = img.shapethresh = img.max()/2.5for x in range(width):for y in range(height):val = round(img[x][y],2) if img[x][y] !=0 else 0ax.annotate(str(val), xy=(y,x),horizontalalignment='center',verticalalignment='center', size=8,color='white' if img[x][y]<thresh else 'black')定义CNN网络架构
import torch.nn as nn
import torch.nn.functional as F# define the CNN architecture
class Net(nn.Module):def __init__(self):super(Net, self).__init__()# convolutional layerself.conv1 = nn.Conv2d(3, 16, 3, padding=1)self.conv2 = nn.Conv2d(16,32,3,padding=1)self.conv3 = nn.Conv2d(32,64,3,padding=1)# max pooling layerself.pool = nn.MaxPool2d(2, 2)self.fc1 = nn.Linear(64*4*4,500)self.fc2 = nn.Linear(500,10)self.dropout = nn.Dropout(0.2)def forward(self, x):# add sequence of convolutional and max pooling layersx = self.pool(F.relu(self.conv1(x)))x = self.pool(F.relu(self.conv2(x)))x = self.pool(F.relu(self.conv3(x)))x = x.view(-1,64*4*4)x = self.dropout(x)x = F.relu(self.fc1(x))x = self.fc2(x)return x# create a complete CNN
model = Net()
print(model)# move tensors to GPU if CUDA is available
if train_on_gpu:model.cuda()指定损失函数和优化器
import torch.optim as optim# specify loss function
criterion = nn.CrossEntropyLoss()# specify optimizer
optimizer = optim.SGD(model.parameters(),lr=0.01)训练网络
# number of epochs to train the model
n_epochs = 30 # you may increase this number to train a final modelvalid_loss_min = np.Inf # track change in validation lossfor epoch in range(1, n_epochs+1):# keep track of training and validation losstrain_loss = 0.0valid_loss = 0.0#################### train the model ####################model.train()for data, target in train_loader:# move tensors to GPU if CUDA is availableif train_on_gpu:data, target = data.cuda(), target.cuda()# clear the gradients of all optimized variablesoptimizer.zero_grad()# forward pass: compute predicted outputs by passing inputs to the modeloutput = model(data)# calculate the batch lossloss = criterion(output, target)# backward pass: compute gradient of the loss with respect to model parametersloss.backward()# perform a single optimization step (parameter update)optimizer.step()# update training losstrain_loss += loss.item()*data.size(0)###################### # validate the model #######################model.eval()for data, target in valid_loader:# move tensors to GPU if CUDA is availableif train_on_gpu:data, target = data.cuda(), target.cuda()# forward pass: compute predicted outputs by passing inputs to the modeloutput = model(data)# calculate the batch lossloss = criterion(output, target)# update average validation loss valid_loss += loss.item()*data.size(0)# calculate average lossestrain_loss = train_loss/len(train_loader.dataset)valid_loss = valid_loss/len(valid_loader.dataset)# print training/validation statistics print('Epoch: {} \tTraining Loss: {:.6f} \tValidation Loss: {:.6f}'.format(epoch, train_loss, valid_loss))# save model if validation loss has decreasedif valid_loss <= valid_loss_min:print('Validation loss decreased ({:.6f} --> {:.6f}). Saving model ...'.format(valid_loss_min,valid_loss))torch.save(model.state_dict(), 'model_cifar.pt')valid_loss_min = valid_loss加载验证损失最低的模型
model.load_state_dict(torch.load('model_cifar.pt'))验证训练的网络
# track test loss
test_loss = 0.0
class_correct = list(0. for i in range(10))
class_total = list(0. for i in range(10))model.eval()
# iterate over test data
for data, target in test_loader:# move tensors to GPU if CUDA is availableif train_on_gpu:data, target = data.cuda(), target.cuda()# forward pass: compute predicted outputs by passing inputs to the modeloutput = model(data)# calculate the batch lossloss = criterion(output, target)# update test loss test_loss += loss.item()*data.size(0)# convert output probabilities to predicted class_, pred = torch.max(output, 1) # compare predictions to true labelcorrect_tensor = pred.eq(target.data.view_as(pred))correct = np.squeeze(correct_tensor.numpy()) if not train_on_gpu else np.squeeze(correct_tensor.cpu().numpy())# calculate test accuracy for each object classfor i in range(batch_size):label = target.data[i]class_correct[label] += correct[i].item()class_total[label] += 1# average test loss
test_loss = test_loss/len(test_loader.dataset)
print('Test Loss: {:.6f}\n'.format(test_loss))for i in range(10):if class_total[i] > 0:print('Test Accuracy of %5s: %2d%% (%2d/%2d)' % (classes[i], 100 * class_correct[i] / class_total[i],np.sum(class_correct[i]), np.sum(class_total[i])))else:print('Test Accuracy of %5s: N/A (no training examples)' % (classes[i]))print('\nTest Accuracy (Overall): %2d%% (%2d/%2d)' % (100. * np.sum(class_correct) / np.sum(class_total),np.sum(class_correct), np.sum(class_total)))可视化测试结果
# obtain one batch of test images
dataiter = iter(test_loader)
images, labels = dataiter.next()
images.numpy()# move model inputs to cuda, if GPU available
if train_on_gpu:images = images.cuda()# get sample outputs
output = model(images)
# convert output probabilities to predicted class
_, preds_tensor = torch.max(output, 1)
preds = np.squeeze(preds_tensor.numpy()) if not train_on_gpu else np.squeeze(preds_tensor.cpu().numpy())# plot the images in the batch, along with predicted and true labels
fig = plt.figure(figsize=(25, 4))
for idx in np.arange(20):ax = fig.add_subplot(2, 20/2, idx+1, xticks=[], yticks=[])imshow(images[idx])ax.set_title("{} ({})".format(classes[preds[idx]], classes[labels[idx]]),color=("green" if preds[idx]==labels[idx].item() else "red"))