Need to change GPU option to CPU in a python pytorch based code

Question

The code basically trains the usual MNIST image dataset but it does the training on a GPU. I need to change this option so the code trains the model using my laptop computer. I need to substitute the .cuda() at the second line for the equivalent in CPU.

I know there are many examples online on how to train neural networks using the MNIST database but what is special about this code is that it does the optimization using a PID controller (commonly used in industry) and I need the code as part of my research.

net = Net(input_size, hidden_size, num_classes)
net.cuda()                                                                 
net.train()                                                                
#Loss and Optimizer
criterion = nn.CrossEntropyLoss()  
optimizer = PIDOptimizer(net.parameters(), lr=learning_rate, weight_decay=0.0001, momentum=0.9, I=I, D=D)
# Train the Model
for epoch in range(num_epochs):
    train_loss_log = AverageMeter()
    train_acc_log = AverageMeter()
    val_loss_log = AverageMeter()
    val_acc_log = AverageMeter()    
    for i, (images, labels) in enumerate(train_loader):  
        # Convert torch tensor to Variable
        images = Variable(images.view(-1, 28*28).cuda())
        labels = Variable(labels.cuda())

Would need to be able to run the code without using the .cuda() option which is for training using a GPU. Need to run it on my PC.

Here's the source code in case needed.

https://github.com/tensorboy/PIDOptimizer

Many thanks, community!

Answer 1

It is better to move up to latest pytorch (1.0.x).

With latest pytorch, it is more easy to manage "device".

Below is a simple example.

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
#Now send existing model to device.
model_ft = model_ft.to(device)
#Now send input to device and so on.
inputs = inputs.to(device)

With this construct, your code automatically uses appropriate device.

Hope this helps!