《深度学习之PyTorch实战计算机视觉》第10章代码 循环神经网络

《深度学习之PyTorch实战计算机视觉》第10章代码

代码

import torch
import torchvision
from torchvision import datasets, transforms
from torch.autograd import Variable
import matplotlib.pyplot as plt
import time
import os
model_path = 'model_name.pth'
model_params_path = 'params_name.pth'

Use_gpu = torch.cuda.is_available()

transform = transforms.Compose(
    [
        transforms.ToTensor(),
        transforms.Normalize(
            mean = [0.5],
            std = [0.5]   
        )
    ]
)


dataset_train = datasets.MNIST(
    root = "./data",
    transform = transform,
    train = True,
    download = True
)

dataset_test = datasets.MNIST(
    root = "./data",
    transform = transform,
    train = False #,
    #download = True #
)

train_load = torch.utils.data.DataLoader(
    dataset = dataset_train,
    batch_size = 64,
    shuffle = True
)


test_load = torch.utils.data.DataLoader(
    dataset = dataset_test,
    batch_size = 64,
    shuffle = True
)




class RNN(torch.nn.Module):

    def __init__(self):
        super(RNN,self).__init__()
        self.rnn = torch.nn.RNN(
            input_size = 28,   #用于指定输入数据的特征数
            hidden_size = 128, #用于指定最后隐藏层的输出特征数
            num_layers = 1,    #用于指定循环层堆叠的数量，默认会使用1
            batch_first = True 
            # 在我们的循环神经网络模型中输入层和输出层用到的数据的默认
            # 维度是(seq, batch, feature)，其中，seq为序列的长度，
            # batch为数据批次的数量，feature为输入或输出的特征数。如果
            # 我们将该参数指定为True，那么输入层和输出层的数据维
            # 度将重新对应为(batch, seq, feature)
        )
        self.output = torch.nn.Linear(128,10)

    def forward(self,input):
        output,_ = self.rnn(input,None)
        output = self.output(output[:,-1,:])
        #因为我们的模型需要处理的是分类问题，所以需要提取
        #最后一个序列的输出结果作为当前循环神经网络模型的输出。
        return output

model = RNN()
if Use_gpu:
    model = model.cuda()

#print(model)

optimizer = torch.optim.Adam(model.parameters())
loss_f = torch.nn.CrossEntropyLoss()


has_been_trained = os.path.isfile(model_path)
if has_been_trained:
    epoch_n = 0
else:
    epoch_n = 5

time_open = time.time()
for epoch in range(epoch_n):
    running_loss = 0.0
    running_correct = 0
    testing_correct = 0
    print("Epoch {}/{}".format(epoch + 1,epoch_n))
    print("-"*20)

    for data in train_load:
        X_train, y_train = data #X_train.shape=torch.Size([64, 1, 28, 28])   y_train.shape=torch.Size([64])
        #print(X_train.shape)
        #print(y_train)
        X_train = X_train.view(-1,28,28)
        if Use_gpu:
            X_train, y_train = Variable(X_train.cuda()), Variable(y_train.cuda())
        else:
            X_train, y_train = Variable(X_train), Variable(y_train)
        y_pred = model(X_train)
        loss = loss_f(y_pred, y_train)
        _, pred = torch.max(y_pred.data, 1)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        running_loss += loss.item()
        running_correct += torch.sum(pred == y_train.data)

    for data in test_load:
        X_test, y_test = data
        X_test = X_test.view(-1,28,28)
        if Use_gpu:
            X_test, y_test = Variable(X_test.cuda()), Variable(y_test.cuda())
        else:
            X_test, y_test = Variable(X_test), Variable(y_test)
        outputs = model(X_test)
        _, pred = torch.max(outputs.data,1)
        testing_correct += torch.sum(pred == y_test.data) 

    print(
            "Loss is:{:.4f}, Train Accuracy is:{:.4f}%, Test Accuracy is:{:.4f}".format(
                running_loss/len(dataset_train),
                100.0*running_correct/len(dataset_train),
                100.0*testing_correct/len(dataset_test)
            )
    )       
time_end = time.time() - time_open
print("程序运行时间:{}分钟{}秒...".format(int(time_end/60),int(time_end)%60))



if has_been_trained:
    model = torch.load(model_path)
else:
    torch.save(model, model_path)



images, label = next(iter(test_load))
images, label = images[0:48], label[0:48]
images_ = images #images.shape = torch.Size([48, 1, 28, 28])
#print(images.shape)
images = images.view(-1,28,28)  #images.shape = torch.Size([48, 28, 28])
#print(images.shape)
if Use_gpu:
    images, label = Variable(images.cuda()), Variable(label.cuda())
else:
    images, label = Variable(images), Variable(label)

pred = model(images)
_,pred = torch.max(pred,1)
print("计算机自动识别出的数字是:\n",[ i.item() for i in pred.data])
print("图片中实际表示的数字是:\n",[ i.item() for i in label])


images_example = torchvision.utils.make_grid(images_) #images_example.shape = torch.Size([3, 182, 242])
#print(images_example.shape)#torch.Size([3, 182, 242])
example = images_example.numpy().transpose(1,2,0)
std = [0.5,0.5,0.5]#std = [0.5,0.5,0.5]
mean = [0.5,0.5,0.5]#mean = [0.5,0.5,0.5]
plt.imshow(example * std + mean)
plt.show() # 6(height) * 8(width)

           

运行效果