Update app.py

app.py

@@ -208,7 +208,9 @@ if __name__ == '__main__':
     gradio_thread.join()
     flask_thread.join()
 '''
+#############################
 
+'''
 from ultralytics import YOLO
 
 # Load a model
@@ -218,3 +220,64 @@ model = YOLO("yolov8n.yaml").load("yolov8n.pt")  # build from YAML and transfer
 
 # Train the model
 results = model.train(data="coco8.yaml", epochs=100, imgsz=640)
+'''
+
+
+import torch  
+import torch.nn as nn  
+import torch.optim as optim  
+from torchvision import datasets, transforms  
+from torch.utils.data import DataLoader  
+  
+# 定义模型  
+class SimpleCNN(nn.Module):  
+    def __init__(self, num_classes=10):  
+        super(SimpleCNN, self).__init__()  
+        self.conv1 = nn.Conv2d(1, 20, 5, 1)  
+        self.conv2 = nn.Conv2d(20, 50, 5, 1)  
+        self.fc1 = nn.Linear(4*4*50, 500)  
+        self.fc2 = nn.Linear(500, num_classes)  
+  
+    def forward(self, x):  
+        x = torch.relu(self.conv1(x))  
+        x = torch.max_pool2d(x, 2, 2)  
+        x = torch.relu(self.conv2(x))  
+        x = torch.max_pool2d(x, 2, 2)  
+        x = x.view(-1, 4*4*50)  
+        x = torch.relu(self.fc1(x))  
+        x = self.fc2(x)  
+        return x  
+  
+# 加载数据集  
+transform = transforms.Compose([  
+    transforms.ToTensor(),  
+    transforms.Normalize((0.5,), (0.5,))  
+])  
+  
+train_dataset = datasets.MNIST(root='./data', train=True, transform=transform, download=True)  
+train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)  
+  
+# 初始化模型和优化器  
+model = SimpleCNN(num_classes=10)  
+optimizer = optim.Adam(model.parameters(), lr=0.001)  
+criterion = nn.CrossEntropyLoss()  
+  
+# 训练模型  
+num_epochs = 5  
+for epoch in range(num_epochs):  
+    for i, (images, labels) in enumerate(train_loader):  
+        # 前向传播  
+        outputs = model(images)  
+        loss = criterion(outputs, labels)  
+  
+        # 反向传播和优化  
+        optimizer.zero_grad()  
+        loss.backward()  
+        optimizer.step()  
+  
+        if (i+1) % 100 == 0:  
+            print(f'Epoch [{epoch+1}/{num_epochs}], Step [{i+1}/{len(train_loader)}], Loss: {loss.item():.4f}')  
+  
+# 保存模型（可选）  
+torch.save(model.state_dict(), 'model.pth')
+