PyTorchRNN

    1. # Example code of a simple RNN, GRU, LSTM on the MNIST dataset.
    3. import torch
    4. import torchvision
    6. # 参数函数，类似激活函数
    7. import torch.nn.functional as F
    9. # standard datasets
    10. import torchvision.datasets as datasets
    12. # 可以对数据集进行转换
    13. import torchvision.transforms as transforms
    15. from torch import optim  # 参数优化
    16. from torch import nn  # 所有神经网络模块
    18. # Gives easier dataset managment by creating mini batches etc
    19. from torch.utils.data import DataLoader
    21. # for nice progress bar!
    22. from tqdm import tqdm
    24. # Set device cuda for GPU if it's available otherwise run on the CPU
    25. device  = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    27. #--------------------------------------------------------------------
    28. # Recurrent neural network (many-to-one)
    29. class RNN(nn.Module):
    30.   def __init__(self,input_size,hidden_size,num_layers,num_classes):
    31.     super(RNN,self).__init__()
    32.     self.hidden_size = hidden_size
    33.     self.num_layers = num_layers
    34.     self.rnn = nn.RNN(input_size,hidden_size,num_layers,batch_first=True)
    35.     self.fc = nn.Linear(hidden_size*seq_len,num_classes)
    36.   def forward(self,x):
    37.     # set initial hidden and cell stats
    38.     h0 = torch.zeros(self.num_layers,x.size(0),self.hidden_size).to(device)
    40.     out,_ = self.rnn(x,h0)
    41.     out = out.reshape(x.shape[0],-1)
    42.     out = self.fc(out)
    43.     return out
    45. #--------------------------------------------------------------------
    46. class RNN_GRU(nn.Module):
    47.   def __int__(self,input_size,hidden_size,num_layers,num_classes):
    48.     super(RNN_GRU,self).__init__()
    49.     self.hidden_size = hidden_size
    50.     self.num_layers = num_layers
    51.     selg.gru = nn.GRU(input_size,hidden_size,num_layers,batch_size=True)
    52.     self.fc = nn.Linear(hidden_size*seq_len,num_classes)
    54.   def forward(self,x):
    55.     h0 = torch.zeros(num_layers,x.size(0),self.hidden_size).to(device)
    57.     out,_ = self.gru(x,h0)
    58.     out = out.reshape(x.size(0),-1)
    59.     out =self.fc(out)
    60.     return out
    62. #--------------------------------------------------------------------
    63. class RNN_LSTM(nn.Module):
    64.   def __init__(self,input_size,hidden_size,num_layers,num_classes):
    65.     super(RNN_LSTM,self).__init__()
    66.     self.hidden_size = hidden_size
    67.     self.num_layers = num_layers
    68.     self.lstm = nn.LSTM(input_size,hidden_size,num_layers,batch_first=True)
    69.     self.fc = nn.Linear(hidden_size*seq_len,num_classes)
    71.   def forward(self,x):
    72.     c0 = torch.zeros(num_layers,x.size(0),self.hidden_size)
    73.     h0 = torch.zeros(num_layers,x.size(0),self.hidden_size)
    75.     out,_ = self.lstm(x,(h0,c0)) 
    76.     # out: tensor of shape (batch_size, seq_length, hidden_size)
    77.     out = out.reshape(out.shape[0], -1)
    78.     out = self.fc(out)
    79.     return out
    81. #--------------------------------------------------------------------
    82. # 参数设置
    83. input_size = 28
    84. num_classes = 10
    85. learning_rate = 0.001
    86. batch_size = 64
    87. num_epochs = 3
    88. num_layers = 2
    89. seq_len = 28
    90. hidden_size = 256
    92. # Load Training and Test data
    93. train_dataset = datasets.MNIST(root='dataset/',train=True,transform=transforms.ToTensor(),download=True)
    94. test_dataset = datasets.MNIST(root='dataset/',train=False,transform=transforms.ToTensor(),download=True)
    95. train_loader = DataLoader(dataset=train_dataset,shuffle=True,batch_size=batch_size)
    96. test_loader = DataLoader(dataset=test_dataset,shuffle=True,batch_size=batch_size)
    98. #--------------------------------------------------------------------
    99. # 初始化网络
    100. model = RNN_LSTM(input_size, hidden_size, num_layers, num_classes).to(device)
    102. # loss and optimizer
    103. criterion= nn.CrossEntropyLoss()
    104. optimizer = optim.Adam(model.parameters(),lr=learning_rate)
    106. # 训练网络
    107. for epoch in range(num_epochs):
    108.     for batch_idxm,(data,targets) in enumerate(tqdm(train_loader)):
    109.         data = data.to(device).squeeze(1)
    110.         targets = targets.to(device)       
    112.         # forward
    113.         outputs = model(data)
    115.         # 计算损失
    116.         loss = criterion(outputs,targets)
    118.         # 向后传播
    119.         loss.backward()
    121.         # 梯度归0
    122.         optimizer.zero_grad()
    124.         # 梯度优化
    125.         optimizer.step()
    127. #--------------------------------------------------------------------
    128. # Check accuracy on training & test to see how good our model
    129. def check_accuracy(loader, model):
    130.     num_correct = 0
    131.     num_samples = 0
    132.     model.eval()
    134.     with torch.no_grad():
    135.         for x, y in loader:
    136.             x = x.to(device=device).squeeze(1)
    137.             y = y.to(device=device)
    139.             outputs = model(x)
    140.             _, indexes = outputs.max(1)
    141.             num_correct += (indexes == y).sum()
    142.             num_samples += indexes.size(0) #  batch_size
    144.     model.train()
    145.     return num_correct/num_samples
    147. #--------------------------------------------------------------------
    148. print(f"Accuracy on training set: {check_accuracy(train_loader, model)*100:.2f}")
    149. print(f"Accuracy on test set: {check_accuracy(test_loader, model)*100:.2f}")
    151. # Accuracy on training set: 10.22
    152. # Accuracy on test set: 10.10