How to use PyTorch GPU acceleration-CNN- MNIST handwritten digit recognition

Hello everyone, I am learning pytroch recently and found this example from the Internet (PyTorch - CNN 卷積神經網絡 - MNIST手寫數字辨識 - HackMD).
I want to try GPU acceleration. I have crawled some information on the forum.
But I still can’t write it, so I would like to ask experienced people to tell me what changes need to be made to the code to achieve GPU computing.

Note: Because I want to visualize the accuracy and error percentage, when I read the article before and said that when I use plt, I need to change the data back to the cpu, so please help me pay attention to it, thank you.

Because I want to visualize the accuracy and error percentage,
when I read the article before and said that when I use plt,
I need to change the data back to the cpu, so please help me pay attention to it, thank you.

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable
from sklearn.model_selection import train_test_split
import matplotlib.pyplot as plt

from keras.datasets import mnist
(X_train, Y_train), (X_test, Y_test) = mnist.load_data()

X_train = X_train.astype('float32') / 255
X_test = X_test.astype('float32') / 255

features_train, features_test, targets_train, targets_test = train_test_split(X_train, Y_train, test_size = 0.2, random_state = 42)

featuresTrain = torch.from_numpy(features_train)
targetsTrain = torch.from_numpy(targets_train).type(torch.LongTensor) # data type is long

featuresTest = torch.from_numpy(features_test)
targetsTest = torch.from_numpy(targets_test).type(torch.LongTensor) # data type is long

# Pytorch train and test TensorDataset
train =,targetsTrain)
test =,targetsTest)

# Hyper Parameters
# batch_size, epoch and iteration
LR = 0.01
batch_size = 100
n_iters = 10000
num_epochs = n_iters / (len(features_train) / batch_size)
num_epochs = int(num_epochs)

# Pytorch DataLoader
train_loader =, batch_size = batch_size, shuffle = True)
test_loader =, batch_size = batch_size, shuffle = True)

# Create CNN Model
class CNN_Model(nn.Module):
    def __init__(self):
        super(CNN_Model, self).__init__()
        # Convolution 1 , input_shape=(1,28,28)
        self.cnn1 = nn.Conv2d(in_channels=1, out_channels=16, kernel_size=5, stride=1, padding=0) #output_shape=(16,24,24)
        self.relu1 = nn.ReLU() # activation
        # Max pool 1
        self.maxpool1 = nn.MaxPool2d(kernel_size=2) #output_shape=(16,12,12)
        # Convolution 2
        self.cnn2 = nn.Conv2d(in_channels=16, out_channels=32, kernel_size=5, stride=1, padding=0) #output_shape=(32,8,8)
        self.relu2 = nn.ReLU() # activation
        # Max pool 2
        self.maxpool2 = nn.MaxPool2d(kernel_size=2) #output_shape=(32,4,4)
        # Fully connected 1 ,#input_shape=(32*4*4)
        self.fc1 = nn.Linear(32 * 4 * 4, 10) 
    def forward(self, x):
        # Convolution 1
        out = self.cnn1(x)
        out = self.relu1(out)
        # Max pool 1
        out = self.maxpool1(out)
        # Convolution 2 
        out = self.cnn2(out)
        out = self.relu2(out)
        # Max pool 2 
        out = self.maxpool2(out)
        out = out.view(out.size(0), -1)
        # Linear function (readout)
        out = self.fc1(out)
        return out

model = CNN_Model()
optimizer = torch.optim.Adam(model.parameters(), lr=LR)   # optimize all cnn parameters
loss_func = nn.CrossEntropyLoss()   # the target label is not one-hotted
input_shape = (-1,1,28,28)

def fit_model(model, loss_func, optimizer, input_shape, num_epochs, train_loader, test_loader):
    # Traning the Model
    #history-like list for store loss & acc value
    training_loss = []
    training_accuracy = []
    validation_loss = []
    validation_accuracy = []
    for epoch in range(num_epochs):
        #training model & store loss & acc / epoch
        correct_train = 0
        total_train = 0
        for i, (images, labels) in enumerate(train_loader):
            # 1.Define variables
            train = Variable(images.view(input_shape))
            labels = Variable(labels)
            # 2.Clear gradients
            # 3.Forward propagation
            outputs = model(train)
            # 4.Calculate softmax and cross entropy loss
            train_loss = loss_func(outputs, labels)
            # 5.Calculate gradients
            # 6.Update parameters
            # 7.Get predictions from the maximum value
            predicted = torch.max(, 1)[1]
            # 8.Total number of labels
            total_train += len(labels)
            # 9.Total correct predictions
            correct_train += (predicted == labels).float().sum() val_acc / epoch
        train_accuracy = 100 * correct_train / float(total_train)
        # loss / epoch

        #evaluate model & store loss & acc / epoch
        correct_test = 0
        total_test = 0
        for images, labels in test_loader:
            # 1.Define variables
            test = Variable(images.view(input_shape))
            # 2.Forward propagation
            outputs = model(test)
            # 3.Calculate softmax and cross entropy loss
            val_loss = loss_func(outputs, labels)
            # 4.Get predictions from the maximum value
            predicted = torch.max(, 1)[1]
            # 5.Total number of labels
            total_test += len(labels)
            # 6.Total correct predictions
            correct_test += (predicted == labels).float().sum() val_acc / epoch
        val_accuracy = 100 * correct_test / float(total_test)
        # val_loss / epoch
        print('Train Epoch: {}/{} Traing_Loss: {} Traing_acc: {:.6f}% Val_Loss: {} Val_accuracy: {:.6f}%'.format(epoch+1, num_epochs,, train_accuracy,, val_accuracy))
    return training_loss, training_accuracy, validation_loss, validation_accuracy

training_loss, training_accuracy, validation_loss, validation_accuracy = fit_model(model, loss_func, optimizer, input_shape, num_epochs, train_loader, test_loader)

# visualization
plt.plot(range(num_epochs), training_loss, 'b-', label='Training_loss')
plt.plot(range(num_epochs), validation_loss, 'g-', label='validation_loss')
plt.title('Training & Validation loss')
plt.xlabel('Number of epochs')
plt.plot(range(num_epochs), training_accuracy, 'b-', label='Training_accuracy')
plt.plot(range(num_epochs), validation_accuracy, 'g-', label='Validation_accuracy')
plt.title('Training & Validation accuracy')
plt.xlabel('Number of epochs')

Hi 浚瑋!

Pytorch tensors can be “moved” to the gpu so that computations
occur – greatly accelerated – on the gpu.

You can created a copy of a cpu tensor that resides on the gpu with:

my_gpu_tensor = my_cpu_tensor.cuda()

If you have a model that is derived from torch.nn.Module, you can
have it move its weights to the gpu with:

my_model = torch.nn.Linear (5, 10)   # Linear is a kind of Module
# my_model is initially on the cpu
my_model.cuda()   # move the model weights to the gpu
# my_model is now on the gpu

If you have a tensor on the gpu – for example the output of your
model that is running on the gpu – you can make a copy of it on
the cpu with:

my_cpu_tensor = my_gpu_tensor.cpu()

Your use of Variable suggests that the example you are working
with is quite old and is written for an old version of pytorch.

(Variable is an old wrapper for Tensor that keeps track of the
information needed for autograd. Newer versions of pytorch – that
are better in lots of ways – have this functionality built directly into
Tensor so Variable is now deprecated.)

Make sure you are using a reasonably up-to-date version of pytorch
and base your learning on up-to-date example code.

The main pytorch web site has a number of (up-to-date) tutorials
that would be a good place to start.

Good luck.

K. Frank

Hi K. Frank
Thank you for informing me of this course information, I will make more use of it.
thank you very much for your reply.

Best regards.