PyTorchCNN

    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
    26. #--------------------------------------------------------------------
    27. class CNN(nn.Module):
    28.     def __init__(self, in_channels=1, num_classes=10):
    29.         super(CNN, self).__init__()
    30.         self.conv1 = nn.Conv1d(in_channels=in_channels,
    31.                     out_channels=8,
    32.                     kernel_size=(3,3),
    33.                     stride=(1,1),
    34.                     padding=(1,1))
    36.         self.pool = nn.MaxPool2d(kernel_size=(2,2),stride=(2,2))
    38.         self.conv2 = nn.Conv2d(in_channels=8,
    39.                     out_channels=16,
    40.                     kernel_size=(3,3),
    41.                     stride=(1,1),
    42.                     padding=(1,1))
    44.         self.fc = nn.Linear(16*7*7, num_classes)
    46.     def forward(self, x):
    47.         x = F.relu(self.conv1(x))
    48.         x = self.pool(x)
    49.         x = F.relu(self.conv2(x))
    50.         x = self.pool(x)
    51.         x = x.reshape(x.shape[0],-1)
    52.         x = self.fc(x)
    53.         return x
    56. #--------------------------------------------------------------------
    57. # Set device cuda for GPU if it's available otherwise run on the CPU
    58. device  = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    60. # 参数设置
    61. in_channels = 1
    62. num_classes = 10
    63. learning_rate = 0.001
    64. batch_size = 64
    65. num_epochs = 3
    67. # Load Training and Test data
    68. train_dataset = datasets.MNIST(root='dataset/',train=True,transform=transforms.ToTensor(),download=True)
    69. test_dataset = datasets.MNIST(root='dataset/',train=False,transform=transforms.ToTensor(),download=True)
    70. train_loader = DataLoader(dataset=train_dataset,shuffle=True,batch_size=batch_size)
    71. test_loader = DataLoader(dataset=test_dataset,shuffle=True,batch_size=batch_size)
    74. #--------------------------------------------------------------------
    75. # 初始化网络
    76. model = CNN(in_channels=in_channels,num_classes=num_classes).to(device)
    78. # loss and optimizer
    79. criterion= nn.CrossEntropyLoss()
    80. optimizer = optim.Adam(model.parameters(),lr=learning_rate)
    82. # 训练网络
    83. for epoch in range(num_epochs):
    84.     for batch_idxm,(data,targets) in enumerate(tqdm(train_loader)):
    85.         data = data.to(device)
    86.         targets = targets.to(device)       
    88.         # forward
    89.         outputs = model(data)
    91.         # 计算损失
    92.         loss = criterion(outputs,targets)
    94.         # 向后传播
    95.         loss.backward()
    97.         # 梯度归0
    98.         optimizer.zero_grad()
    100.         # 梯度优化
    101.         optimizer.step()
    104. #--------------------------------------------------------------------
    105. # Check accuracy on training & test to see how good our model
    106. def check_accuracy(loader, model):
    107.     num_correct = 0
    108.     num_samples = 0
    109.     model.eval()
    111.     with torch.no_grad():
    112.         for x, y in loader:
    113.             x = x.to(device=device)
    114.             y = y.to(device=device)
    116.             outputs = model(x)
    117.             _, indexes = outputs.max(1)
    118.             num_correct += (indexes == y).sum()
    119.             num_samples += indexes.size(0) #  batch_size
    121.     model.train()
    122.     return num_correct/num_samples
    125. #--------------------------------------------------------------------
    126. print(f"Accuracy on training set: {check_accuracy(train_loader, model)*100:.2f}")
    127. print(f"Accuracy on test set: {check_accuracy(test_loader, model)*100:.2f}")
    129. # Accuracy on training set: 9.62
    130. # Accuracy on test set: 9.59