修改requirements.txt文件，添加main函数判断

Lab1/code/3.2.py

@@ -20,77 +20,78 @@ class Model_3_2(nn.Module):
         return x
 
 
-learning_rate = 5e-2
-num_epochs = 10
-batch_size = 4096
-num_classes = 10
-device = "cuda:0" if torch.cuda.is_available() else "cpu"
+if __name__ == "__main__":
+    learning_rate = 5e-2
+    num_epochs = 10
+    batch_size = 4096
+    num_classes = 10
+    device = "cuda:0" if torch.cuda.is_available() else "cpu"
 
-transform = transforms.Compose(
-    [
-        transforms.ToTensor(),
-        transforms.Normalize((0.5,), (1.0,)),
-    ]
-)
-train_dataset = datasets.FashionMNIST(
-    root="../dataset", train=True, transform=transform, download=True
-)
-test_dataset = datasets.FashionMNIST(
-    root="../dataset", train=False, transform=transform, download=True
-)
-train_loader = DataLoader(
-    dataset=train_dataset,
-    batch_size=batch_size,
-    shuffle=True,
-    num_workers=14,
-    pin_memory=True,
-)
-test_loader = DataLoader(
-    dataset=test_dataset,
-    batch_size=batch_size,
-    shuffle=True,
-    num_workers=14,
-    pin_memory=True,
-)
-
-model = Model_3_2(num_classes).to(device)
-criterion = nn.CrossEntropyLoss()
-optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)
-
-for epoch in range(num_epochs):
-    total_epoch_loss = 0
-    model.train()
-    for index, (images, targets) in tqdm(
-        enumerate(train_loader), total=len(train_loader)
-    ):
-        optimizer.zero_grad()
-
-        images = images.to(device)
-        targets = targets.to(device)
-
-        one_hot_targets = (
-            torch.nn.functional.one_hot(targets, num_classes=num_classes)
-            .to(device)
-            .to(dtype=torch.float32)
-        )
-
-        outputs = model(images)
-        loss = criterion(outputs, one_hot_targets)
-        total_epoch_loss += loss
-
-        loss.backward()
-        optimizer.step()
-
-    model.eval()
-    total_acc = 0
-    with torch.no_grad():
-        for index, (image, targets) in tqdm(
-            enumerate(test_loader), total=len(test_loader)
-        ):
-            image = image.to(device)
-            targets = targets.to(device)
-            outputs = model(image)
-            total_acc += (outputs.argmax(1) == targets).sum()
-    print(
-        f"Epoch {epoch + 1}/{num_epochs} Train, Loss: {total_epoch_loss}, Acc: {total_acc / len(test_dataset)}"
+    transform = transforms.Compose(
+        [
+            transforms.ToTensor(),
+            transforms.Normalize((0.5,), (1.0,)),
+        ]
     )
+    train_dataset = datasets.FashionMNIST(
+        root="../dataset", train=True, transform=transform, download=True
+    )
+    test_dataset = datasets.FashionMNIST(
+        root="../dataset", train=False, transform=transform, download=True
+    )
+    train_loader = DataLoader(
+        dataset=train_dataset,
+        batch_size=batch_size,
+        shuffle=True,
+        num_workers=14,
+        pin_memory=True,
+    )
+    test_loader = DataLoader(
+        dataset=test_dataset,
+        batch_size=batch_size,
+        shuffle=True,
+        num_workers=14,
+        pin_memory=True,
+    )
+
+    model = Model_3_2(num_classes).to(device)
+    criterion = nn.CrossEntropyLoss()
+    optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)
+
+    for epoch in range(num_epochs):
+        total_epoch_loss = 0
+        model.train()
+        for index, (images, targets) in tqdm(
+            enumerate(train_loader), total=len(train_loader)
+        ):
+            optimizer.zero_grad()
+
+            images = images.to(device)
+            targets = targets.to(device)
+
+            one_hot_targets = (
+                torch.nn.functional.one_hot(targets, num_classes=num_classes)
+                .to(device)
+                .to(dtype=torch.float32)
+            )
+
+            outputs = model(images)
+            loss = criterion(outputs, one_hot_targets)
+            total_epoch_loss += loss
+
+            loss.backward()
+            optimizer.step()
+
+        model.eval()
+        total_acc = 0
+        with torch.no_grad():
+            for index, (image, targets) in tqdm(
+                enumerate(test_loader), total=len(test_loader)
+            ):
+                image = image.to(device)
+                targets = targets.to(device)
+                outputs = model(image)
+                total_acc += (outputs.argmax(1) == targets).sum()
+        print(
+            f"Epoch {epoch + 1}/{num_epochs} Train, Loss: {total_epoch_loss}, Acc: {total_acc / len(test_dataset)}"
+        )