# 构建MyDataset实例
transform = transforms.Compose([transforms.ToTensor()])
train_data = TrafficDataSet(TrainDate, transform)
test_data = TrafficDataSet(TestDate,transform)
# 构建DataLoder
train_loader = DataLoader(dataset=train_data, batch_size=BATCH_SIZE, shuffle=True)
test_loader = DataLoader(dataset=test_data, batch_size=BATCH_SIZE, shuffle=True)
# ============================ step 2/5 模型 ============================
net = LeNet(classes=12)
net.initialize_weights()
# ============================ step 3/5 损失函数 ============================
criterion = nn.CrossEntropyLoss() # 选择损失函数
# ============================ step 4/5 优化器 ============================
optimizer = optim.SGD(net.parameters(), lr=LR, momentum=0.9) # 选择优化器
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.1) # 设置学习率下降策略
# ============================ step 5/5 训练集训练 ============================
for epoch in range(MAX_EPOCH):
print("-------- start train Epoch [{:0>3}] ---------".format(epoch+1))
loss_mean = 0.
correct = 0.
