上传文件至「/」

Learning_stage1.py

@@ -0,0 +1,78 @@
+import torch
+from torch import nn
+from torch.utils.data import DataLoader
+from torchvision import datasets
+from torchvision.transforms import v2
+import datetime
+
+time = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print(device)
+transform = v2.Compose([v2.ToImage() , v2.ToDtype(torch.float32, scale=True)])
+# Download training data from open datasets.
+training_data = datasets.FashionMNIST(
+    root="data",
+    train=True,
+    download=True,
+    transform=transform,
+)
+
+# Download test data from open datasets.
+test_data = datasets.FashionMNIST(
+    root="data",
+    train=False,
+    download=True,
+    transform=transform,
+)
+
+train_loader = DataLoader(training_data, batch_size=64, shuffle=True)
+test_loader = DataLoader(test_data, batch_size=64, shuffle=False)
+
+class NeuralNetwork(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.flatten = nn.Flatten()
+        self.network = nn.Sequential(
+            nn.Linear(784, 512),
+            nn.ReLU(),
+            nn.Linear(512, 512),
+            nn.ReLU(),
+            nn.Linear(512, 10),
+        )
+    def forward(self, x):
+        x = self.flatten(x)
+        return self.network(x)
+
+model = NeuralNetwork().to(device)
+print(model)
+loss_fn = nn.CrossEntropyLoss()
+optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
+
+epochs = 10
+for epoch in range(epochs):
+    model.train()
+    for X, y in train_loader:
+        X = X.to(device)
+        y = y.to(device)
+        pred = model(X)
+        loss = loss_fn(pred, y)
+        loss.backward()
+        optimizer.step()
+        optimizer.zero_grad()
+    print("Epoch:", epoch+1,"Loss:", loss.item())
+
+model.eval()
+correct = 0
+
+with torch.no_grad():
+        for X, y in test_loader:
+            X = X.to(device)
+            y = y.to(device)
+            pred = model(X)
+            correct += (pred.argmax(1) == y).float().sum().item()
+print("Accuracy:", correct/len(test_loader))
+
+torch.save(model.state_dict(), f"fashion_mnist_epochs{epochs}_{time}.pth")
+
+print("Model saved")