173 lines
5.5 KiB
Python
173 lines
5.5 KiB
Python
import numpy as np
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from matplotlib import pyplot as plt
|
|
from torch.utils.data import DataLoader
|
|
from torchvision import datasets, transforms
|
|
|
|
# Super parameters ------------------------------------------------------------------------------------
|
|
batch_size = 64
|
|
learning_rate = 0.01
|
|
momentum = 0.5
|
|
EPOCH = 10
|
|
|
|
# Prepare dataset ------------------------------------------------------------------------------------
|
|
transform = transforms.Compose(
|
|
[transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]
|
|
)
|
|
|
|
train_dataset = datasets.MNIST(
|
|
root="./data/mnist", train=True, download=True, transform=transform
|
|
)
|
|
test_dataset = datasets.MNIST(
|
|
root="./data/mnist", train=False, download=True, transform=transform
|
|
)
|
|
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
|
|
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)
|
|
|
|
|
|
# Design model using class ------------------------------------------------------------------------------
|
|
class Net(torch.nn.Module):
|
|
def __init__(self):
|
|
super(Net, self).__init__()
|
|
self.fc1 = torch.nn.Linear(
|
|
28 * 28, 512
|
|
) # Flatten 28x28 input to a vector of 512 units
|
|
self.fc2 = torch.nn.Linear(512, 128) # First hidden layer
|
|
self.dropout = torch.nn.Dropout(p=0.5) # 添加 Dropout 层，丢弃概率为 0.5
|
|
self.fc3 = torch.nn.Linear(
|
|
128, 10
|
|
) # Output layer with 10 units (one for each digit)
|
|
|
|
def forward(self, x):
|
|
x = x.view(-1, 28 * 28) # Flatten the image
|
|
x = self.fc1(x) # the first linear layer
|
|
x = F.relu(x) # Activation
|
|
x = self.fc2(x) # the second linear layer
|
|
x = F.relu(x) # Activation
|
|
x = self.dropout(x) # 在最后一个全连接层之前应用 Dropout
|
|
x = self.fc3(x) # No activation on the final output layer (for classification)
|
|
return x
|
|
|
|
|
|
model = Net()
|
|
|
|
# Construct loss and optimizer ----------------------------------------------------------------------
|
|
criterion = (
|
|
torch.nn.CrossEntropyLoss()
|
|
) # CrossEntropy loss for multi-class classification
|
|
optimizer = torch.optim.SGD(
|
|
model.parameters(), lr=learning_rate, momentum=momentum, weight_decay=0.01
|
|
) # SGD optimizer
|
|
|
|
|
|
# Train and Test CLASS -----------------------------------------------------------------------------------
|
|
def train(epoch):
|
|
running_loss = 0.0
|
|
running_total = 0
|
|
running_correct = 0
|
|
for batch_idx, data in enumerate(train_loader):
|
|
inputs, target = data
|
|
optimizer.zero_grad()
|
|
|
|
# forward + backward + update
|
|
outputs = model(inputs)
|
|
loss = criterion(outputs, target)
|
|
|
|
loss.backward()
|
|
optimizer.step()
|
|
|
|
# Accumulate loss and accuracy
|
|
running_loss += loss.item()
|
|
_, predicted = torch.max(outputs.data, dim=1)
|
|
running_total += target.shape[0]
|
|
running_correct += (predicted == target).sum().item()
|
|
|
|
# Print after each epoch
|
|
avg_loss = running_loss / len(train_loader)
|
|
avg_acc = 100 * running_correct / running_total
|
|
print(
|
|
"[%d / %d]: Training Loss: %.3f, Training Accuracy: %.2f %%"
|
|
% (epoch + 1, EPOCH, avg_loss, avg_acc)
|
|
)
|
|
|
|
|
|
# ... existing code ...
|
|
|
|
|
|
def add_noise(images, noise_level=0.5):
|
|
"""
|
|
添加噪声函数
|
|
参数:
|
|
images: 输入数据 batch (torch.Tensor)
|
|
noise_level: 噪声强度
|
|
返回:
|
|
添加噪声后的数据
|
|
"""
|
|
noise = noise_level * torch.randn_like(images) # 生成随机噪声
|
|
noisy_images = images + noise # 将噪声添加到数据中
|
|
noisy_images = torch.clamp(noisy_images, 0.0, 1.0) # 限制值范围在 [0, 1]
|
|
return noisy_images
|
|
|
|
|
|
def test(epoch, noise_level=0.2):
|
|
correct = 0
|
|
total = 0
|
|
with torch.no_grad():
|
|
for data in test_loader:
|
|
images, labels = data
|
|
# 添加噪声到测试数据
|
|
noisy_images = add_noise(images, noise_level=noise_level)
|
|
outputs = model(noisy_images)
|
|
_, predicted = torch.max(outputs.data, dim=1)
|
|
total += labels.size(0)
|
|
correct += (predicted == labels).sum().item()
|
|
acc = correct / total
|
|
print(
|
|
"[%d / %d]: Accuracy on test set with noise level %.2f after epoch %d: %.1f %%"
|
|
% (epoch + 1, EPOCH, noise_level, epoch + 1, 100 * acc)
|
|
)
|
|
return 100 * acc
|
|
|
|
|
|
# Add EarlyStopping class after imports
|
|
class EarlyStopping:
|
|
def __init__(self, patience=3, min_delta=0):
|
|
self.patience = patience
|
|
self.min_delta = min_delta
|
|
self.counter = 0
|
|
self.best_acc = None
|
|
self.should_stop = False
|
|
|
|
def __call__(self, acc):
|
|
if self.best_acc is None:
|
|
self.best_acc = acc
|
|
elif acc <= self.best_acc + self.min_delta:
|
|
self.counter += 1
|
|
if self.counter >= self.patience:
|
|
self.should_stop = True
|
|
else:
|
|
self.best_acc = acc
|
|
self.counter = 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
acc_list_test_noisy = []
|
|
early_stopping = EarlyStopping(patience=3)
|
|
|
|
for epoch in range(EPOCH):
|
|
train(epoch)
|
|
acc_test_noisy = test(epoch, noise_level=0.8)
|
|
acc_list_test_noisy.append(acc_test_noisy)
|
|
|
|
# Early stopping check
|
|
early_stopping(acc_test_noisy)
|
|
if early_stopping.should_stop:
|
|
print(f"Early stopping triggered at epoch {epoch + 1}")
|
|
break
|
|
|
|
plt.plot(range(1, len(acc_list_test_noisy) + 1), acc_list_test_noisy)
|
|
plt.xlabel("Epoch")
|
|
plt.ylabel("Accuracy On Noisy Test Set")
|
|
plt.show()
|