PyTorch是一个开源的机器学习库，它为深度学习应用提供了丰富的API和工具。然而，将PyTorch模型转换为独立可执行文件（EXE）可能会有很多好处，例如便于部署、加速实际应用等。本教程将向您说明如何将PyTorch模型转换为EXE文件。
### 需要安装的库
- PyTorch
- torchvision
- torchsummary
- 依赖库：NumPy, Matplotlib
- PyInstaller
请确保已通过`pip`安装了这些库。
### 1.创建用于测试的PyTorch模型
在此教程中，我们使用一个简单的卷积神经网络（CNN）鉴别MNIST数据集中的手写数字。
创建一个新的Python文件，命名为`mnist_cnn.py`，并输入以下代码：
```python
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
import torchsummary
# 1. Define the CNN Model
class MnistCNN(nn.Module):
def __init__(self):
super(MnistCNN, self).__init__()
self.conv1 = nn.Conv2d(1, 32, 3)
self.pool1 = nn.MaxPool2d(2, 2)
self.conv2 = nn.Conv2d(32, 64, 3)
self.pool2 = nn.MaxPool2d(2, 2)
self.dropout1 = nn.Dropout(0.25)
self.fc1 = nn.Linear(64 * 6 * 6, 128)
self.dropout2 = nn.Dropout(0.5)
self.fc2 = nn.Linear(128, 10)
def forward(self, x):
x = self.pool1(F.relu(self.conv1(x)))
x = self.pool2(F.relu(self.conv2(x)))
x = x.view(-1, 64 * 6 * 6)
x = self.dropout1(x)
x = F.relu(self.fc1(x))
x = self.dropout2(x)
x = self.fc2(x)
return x
# 2. Load the MNIST dataset
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))])
trainset = torchvision.datasets.MNIST(root='./data', train=True,
download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=100,
shuffle=True, num_workers=2)
testset = torchvision.datasets.MNIST(root='./data', train=False,
download=True, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=100,
shuffle=False, num_workers=2)
# 3. Train the model
def train(model, device, train_loader, optimizer, epoch):
model.train()
for batch_idx, (data, target) in enumerate(train_loader):
data, target = data.to(device), target.to(device)
optimizer.zero_grad()
output = model(data)
loss = F.cross_entropy(output, target)
loss.backward()
optimizer.step()
def main():
model = MnistCNN().to("cpu")
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
for epoch in range(1, 6):
print("Epoch: {}".format(epoch))
train(model, "cpu", trainloader, optimizer, epoch)
torch.save(model.state_dict(), "mnist_cnn.pth")
if __name__ == '__main__':
main()
```
这段代码定义了一个简单的CNN模型，加载MNIST数据集并对模型进行了训练。如果你想要跳过训练步骤，直接使用预训练好的模型。你也可以使用你自己的模型以及权重。
### 2.创建一个用于预测的脚本
接下来，我们需要创建一个脚本来加载训练好的权重并运行预测。在本例中，我们使用单图像进行预测。创建一个新的Python文件，命名为`predict.py`，并输入以下代码：
```python
import torch
import torch.nn as nn
from mnist_cnn import MnistCNN
import numpy as np
from PIL import Image
def predict(image_file):
model = MnistCNN()
model.load_state_dict(torch.load("mnist_cnn.pth", map_location=torch.device("cpu")))
model.eval()
img = Image.open(image_file)
img = img.resize((28, 28)).convert("L")
img = np.asarray(img, dtype=np.float32)[np.newaxis, np.newaxis, :]
img = (img - 0.1307) / 0.3081
output = model(torch.from_numpy(img))
_, predicted = torch.max(output, 1)
print("Predicted Label: ", predicted.item())
if __name__ == "__main__":
import sys
if len(sys.argv) != 2:
print("USAGE: python predict.py path/to/image")
else:
predict(sys.argv[1])
```
此脚本会使用`mnist_cnn.py`中定义的CNN模型。现在我们可以使用预测脚本来识别MNIST数据集中的图像。
### 3.将PyTorch模型转换为EXE文件
接下来，我们将使用PyInstaller将预测脚本转换为可执行文件。请确保已通过`pip`安装了PyInstaller。
打开终端，输入如下命令：
```
pyinstaller predict.py --onefile
```
PyInstaller会将所有依赖项捆绑在一个可执行文件中。生成的EXE文件可以在`dist`目录下找到。
### 4.使用EXE文件进行预测
最后，将训练好的权重文件（`mnist_cnn.pth`）复制到生成的可执行文件的目录下。现在，您可以通过以下命令使用EXE文件进行预测：
```
./predict.exe path/to/image
```
这就是如何将PyTorch模型转换为EXE文件的方法。注意，生成的EXE文件可能会比你预期的更大，因为它包含了所有依赖库。如有需要，可以使用UPX等工具进行文件压缩。
