需要在基于python pytorch的代码中将GPU选项更改为CPU

时间:2019-02-06 00:16:04

标签: deep-learning gpu pytorch

该代码基本上训练了普通的MNIST图像数据集,但是它在GPU上进行了训练。我需要更改此选项,以便代码使用我的便携式计算机训练模型。我需要在第二行用.cuda()替换CPU中的等效值。

我知道在线上有很多关于如何使用MNIST数据库训练神经网络的示例,但是此代码的特殊之处在于它使用PID控制器(行业中常用)进行优化,因此我需要代码作为一部分我的研究。

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())

将需要能够在不使用用于训练GPU的.cuda()选项的情况下运行代码。需要在我的PC上运行。

如果需要,这里是源代码。

https://github.com/tensorboy/PIDOptimizer

非常感谢,社区!

1 个答案:

答案 0 :(得分:2)

最好升级到最新的pytorch(1.0.x)。

使用最新的pytorch,可以更轻松地管理“设备”。

下面是一个简单的例子。

IOException

使用这种结构,您的代码会自动使用适当的设备。

希望这会有所帮助!