PyTorch

    1. #  用pytorch简单的构建全连接层神经网络
    2. #  对手写字进行识别
    4. import torch
    5. import torchvision
    7. # 参数函数，类似激活函数
    8. import torch.nn.functional as F
    10. # standard datasets
    11. import torchvision.datasets as datasets
    13. # 可以对数据集进行转换
    14. import torchvision.transforms as transforms
    16. from torch import optim  # 参数优化
    17. from torch import nn  # 所有神经网络模块
    19. # Gives easier dataset managment by creating mini batches etc
    20. from torch.utils.data import DataLoader
    22. # for nice progress bar!
    23. from tqdm import tqdm
    25. #--------------------------------------------------------------------
    26. class NN(nn.Module):
    27.     def __init__(self, input_size, num_classes):
    28.         super(NN, self).__init__()
    29.         self.fc1 = nn.Linear(input_size, 50)
    30.         self.fc2 = nn.Linear(50, num_classes)
    32.     def forward(self, x):
    33.         x = F.relu(self.fc1(x))
    34.         x = self.fc2(x)
    35.         return x
    37. #--------------------------------------------------------------------
    38. # Set device cuda for GPU if it's available otherwise run on the CPU
    39. device  = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    41. # 参数设置
    42. input_size = 784
    43. num_classes = 10
    44. learning_rate = 0.001
    45. batch_size = 64
    46. num_epochs = 3
    48. #--------------------------------------------------------------------
    49. # Load Training and Test data
    50. train_dataset = datasets.MNIST(root='dataset/',train=True,transform=transforms.ToTensor(),download=True)
    51. test_dataset = datasets.MNIST(root='dataset/',train=False,transform=transforms.ToTensor(),download=True)
    52. train_loader = DataLoader(dataset=train_dataset,shuffle=True,batch_size=batch_size)
    53. test_loader = DataLoader(dataset=test_dataset,shuffle=True,batch_size=batch_size)
    55. #--------------------------------------------------------------------
    56. # 初始化网络
    57. model = NN(input_size=input_size,num_classes=num_classes).to(device)
    59. # loss and optimizer
    60. criterion= nn.CrossEntropyLoss()
    61. optimizer = optim.Adam(model.parameters(),lr=learning_rate)
    63. # 训练网络
    64. for epoch in range(num_epochs):
    65.     for batch_idxm,(data,targets) in enumerate(tqdm(train_loader)):
    66.         data = data.to(device)
    67.         targets = targets.to(device)
    69.         # [64,1,28,28] -> [64,1*28*28]
    70.         print(data.shape)
    71.         data = data.reshape(data.shape[0],-1)
    73.         # 计算损失
    74.         outputs = model(data)
    75.         loss = criterion(outputs,targets)
    77.         # 向后传播
    78.         loss.backward()
    80.         # 梯度归0
    81.         optimizer.zero_grad()
    83.         # 梯度优化
    84.         optimizer.step()
    87. #--------------------------------------------------------------------
    88. # model save and load
    89. def save_checkpoint(state, filename="my_checkpoint.pth.tar"):
    90.     print("=> Saving checkpoint")
    91.     torch.save(state, filename)
    94. def load_checkpoint(checkpoint, model, optimizer):
    95.     print("=> Loading checkpoint")
    96.     model.load_state_dict(checkpoint["state_dict"])
    97.     optimizer.load_state_dict(checkpoint["optimizer"])
    99. checkpoint = {"state_dict": model.state_dict(), "optimizer": optimizer.state_dict()}
    100. # Try save checkpoint
    101. save_checkpoint(checkpoint)
    103. # Try load checkpoint
    104. load_checkpoint(torch.load("my_checkpoint.pth.tar"), model, optimizer)
    107. #--------------------------------------------------------------------
    108. # Check accuracy on training & test to see how good our model
    109. def check_accuracy(loader, model):
    110.     num_correct = 0
    111.     num_samples = 0
    112.     model.eval()
    114.     with torch.no_grad():
    115.         for x, y in loader:
    116.             x = x.to(device=device)
    117.             y = y.to(device=device)
    118.             x = x.reshape(x.shape[0], -1)
    120.             outputs = model(x)
    121.             _, indexes = outputs.max(1)
    122.             num_correct += (indexes == y).sum()
    123.             num_samples += indexes.size(0) #  batch_size
    125.     model.train()
    126.     return num_correct/num_samples
    129. #--------------------------------------------------------------------
    130. print(f"Accuracy on training set: {check_accuracy(train_loader, model)*100:.2f}")
    131. print(f"Accuracy on test set: {check_accuracy(test_loader, model)*100:.2f}")
    133. # Accuracy on training set: 15.65
    134. # Accuracy on test set: 15.53