怎么用GPU运行python代码
时间: 2023-08-25 12:07:36 浏览: 325
使用GPU运行Python代码需要安装CUDA和cuDNN等GPU相关的库,并且需要安装支持GPU的Python库,比如TensorFlow和PyTorch等。
以下是使用TensorFlow和PyTorch在GPU上运行Python代码的示例:
1. 使用TensorFlow在GPU上运行Python代码
```python
import tensorflow as tf
# 指定使用GPU
physical_devices = tf.config.list_physical_devices('GPU')
tf.config.experimental.set_memory_growth(physical_devices[0], True)
# 定义模型
model = tf.keras.Sequential([
tf.keras.layers.Dense(32, activation='relu', input_shape=(784,)),
tf.keras.layers.Dense(10, activation='softmax')
])
# 编译模型
model.compile(optimizer='adam',
loss='sparse_categorical_crossentropy',
metrics=['accuracy'])
# 加载数据
(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()
# 数据预处理
x_train, x_test = x_train / 255.0, x_test / 255.0
# 训练模型
model.fit(x_train, y_train, epochs=5, validation_data=(x_test, y_test))
```
2. 使用PyTorch在GPU上运行Python代码
```python
import torch
# 指定使用GPU
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# 定义模型
model = torch.nn.Sequential(
torch.nn.Linear(784, 32),
torch.nn.ReLU(),
torch.nn.Linear(32, 10),
torch.nn.Softmax(dim=1)
)
# 将模型放到GPU上
model.to(device)
# 定义损失函数和优化器
criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters())
# 加载数据
train_dataset = torchvision.datasets.MNIST(root='./data', train=True, transform=torchvision.transforms.ToTensor(), download=True)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=64, shuffle=True)
# 训练模型
for epoch in range(5):
for i, (inputs, labels) in enumerate(train_loader):
# 将数据放到GPU上
inputs, labels = inputs.to(device), labels.to(device)
# 前向传播
outputs = model(inputs)
loss = criterion(outputs, labels)
# 反向传播和优化
optimizer.zero_grad()
loss.backward()
optimizer.step()
```
需要注意的是,使用GPU运行Python代码需要确保你的电脑上有支持GPU的硬件设备(例如NVIDIA的显卡),并且安装了相应的驱动程序。
阅读全文