commit 74121c3bfa472881d6037a55ccd28b639fc97c04 Author: guzhujushi Date: Tue Aug 25 23:15:10 2026 +0800 上传文件至「/」 diff --git a/Learning_stage1.py b/Learning_stage1.py new file mode 100644 index 0000000..992bb17 --- /dev/null +++ b/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")