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

Lab1/code/2.2.py

@@ -38,56 +38,57 @@ class My_Dataset(Dataset):
         return x, y
 
 
-learning_rate = 5e-2
-num_epochs = 10
-batch_size = 1024
-device = "cuda:0" if torch.cuda.is_available() else "cpu"
+if __name__ == "__main__":
+    learning_rate = 5e-2
+    num_epochs = 10
+    batch_size = 1024
+    device = "cuda:0" if torch.cuda.is_available() else "cpu"
 
-dataset = My_Dataset()
-dataloader = DataLoader(
-    dataset=dataset,
-    batch_size=batch_size,
-    shuffle=True,
-    num_workers=14,
-    pin_memory=True,
-)
-
-model = Model_2_2().to(device)
-criterion = nn.BCELoss()
-optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
-
-for epoch in range(num_epochs):
-    total_epoch_loss = 0
-    total_epoch_pred = 0
-    total_epoch_target = 0
-    for index, (x, targets) in tqdm(enumerate(dataloader), total=len(dataloader)):
-        optimizer.zero_grad()
-
-        x = x.to(device)
-        targets = targets.to(device)
-
-        x = x.unsqueeze(1)
-        targets = targets.unsqueeze(1)
-        y_pred = model(x)
-        loss = criterion(y_pred, targets)
-        total_epoch_loss += loss.item()
-        total_epoch_target += targets.sum().item()
-        total_epoch_pred += y_pred.sum().item()
-
-        loss.backward()
-        optimizer.step()
-
-    print(
-        f"Epoch {epoch + 1}/{num_epochs}, Loss: {total_epoch_loss}, Acc: {1 - abs(total_epoch_pred - total_epoch_target) / total_epoch_target}"
+    dataset = My_Dataset()
+    dataloader = DataLoader(
+        dataset=dataset,
+        batch_size=batch_size,
+        shuffle=True,
+        num_workers=14,
+        pin_memory=True,
     )
 
-with torch.no_grad():
-    test_data = (np.array([[2]]) - dataset.min_x) / (dataset.max_x - dataset.min_x)
-    test_data = Variable(
-        torch.tensor(test_data, dtype=torch.float64), requires_grad=False
-    ).to(device)
-    predicted = model(test_data).to("cpu")
-    print(
-        f"Model weights: {model.linear.weight.item()}, bias: {model.linear.bias.item()}"
-    )
-    print(f"Prediction for test data: {predicted.item()}")
+    model = Model_2_2().to(device)
+    criterion = nn.BCELoss()
+    optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
+
+    for epoch in range(num_epochs):
+        total_epoch_loss = 0
+        total_epoch_pred = 0
+        total_epoch_target = 0
+        for index, (x, targets) in tqdm(enumerate(dataloader), total=len(dataloader)):
+            optimizer.zero_grad()
+
+            x = x.to(device)
+            targets = targets.to(device)
+
+            x = x.unsqueeze(1)
+            targets = targets.unsqueeze(1)
+            y_pred = model(x)
+            loss = criterion(y_pred, targets)
+            total_epoch_loss += loss.item()
+            total_epoch_target += targets.sum().item()
+            total_epoch_pred += y_pred.sum().item()
+
+            loss.backward()
+            optimizer.step()
+
+        print(
+            f"Epoch {epoch + 1}/{num_epochs}, Loss: {total_epoch_loss}, Acc: {1 - abs(total_epoch_pred - total_epoch_target) / total_epoch_target}"
+        )
+
+    with torch.no_grad():
+        test_data = (np.array([[2]]) - dataset.min_x) / (dataset.max_x - dataset.min_x)
+        test_data = Variable(
+            torch.tensor(test_data, dtype=torch.float64), requires_grad=False
+        ).to(device)
+        predicted = model(test_data).to("cpu")
+        print(
+            f"Model weights: {model.linear.weight.item()}, bias: {model.linear.bias.item()}"
+        )
+        print(f"Prediction for test data: {predicted.item()}")