Day40打卡 @浙大疏锦行
知识点回顾:
- 彩色和灰度图片测试和训练的规范写法:封装在函数中
- 展平操作:除第一个维度batchsize外全部展平
- 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)
@浙大疏锦行