Day40打卡 @浙大疏锦行

知识点回顾：

1. 彩色和灰度图片测试和训练的规范写法：封装在函数中
2. 展平操作：除第一个维度batchsize外全部展平
3. dropout操作：训练阶段随机丢弃神经元，测试阶段eval模式关闭dropout

import torchdef flatten_tensor(x):"""保留batch维度，展平其余所有维度"""return x.view(x.size(0), -1)

class NeuralNetwork(torch.nn.Module):def __init__(self, input_dim, hidden_dim, dropout_prob=0.5):super().__init__()self.layer1 = torch.nn.Linear(input_dim, hidden_dim)self.dropout = torch.nn.Dropout(dropout_prob)self.layer2 = torch.nn.Linear(hidden_dim, 10)def forward(self, x):x = torch.relu(self.layer1(x))x = self.dropout(x)  # 训练时激活，测试时自动关闭return self.layer2(x)

model = NeuralNetwork(input_dim=784, hidden_dim=256)# 训练阶段
model.train()
output_train = model(flattened_grayscale)  # Dropout生效# 测试阶段
model.eval()
with torch.no_grad():output_test = model(flattened_grayscale)  # Dropout自动关闭

def train_epoch(model, dataloader, criterion, optimizer):model.train()for inputs, targets in dataloader:inputs = flatten_tensor(inputs)optimizer.zero_grad()outputs = model(inputs)loss = criterion(outputs, targets)loss.backward()optimizer.step()def evaluate(model, dataloader):model.eval()total_correct = 0with torch.no_grad():for inputs, targets in dataloader:inputs = flatten_tensor(inputs)outputs = model(inputs)_, predicted = torch.max(outputs, 1)total_correct += (predicted == targets).sum().item()return total_correct / len(dataloader.dataset)

@浙大疏锦行​​​​​​​