From 6713993f4b14e2f36d477840ce1455baa9fac4f4 Mon Sep 17 00:00:00 2001 From: carry <2641257231@qq.com> Date: Thu, 3 Aug 2023 22:57:10 +0800 Subject: [PATCH] =?UTF-8?q?=E6=9B=B4=E6=96=B0=E4=BA=86=E6=A8=A1=E5=9E=8B?= =?UTF-8?q?=E4=BF=9D=E5=AD=98=E7=9A=84=E6=9D=A1=E4=BB=B6=EF=BC=8C=E4=BD=BF?= =?UTF-8?q?=E7=94=A8=E4=BA=86ReduceLROnPlateau=E5=AD=A6=E4=B9=A0=E7=8E=87?= =?UTF-8?q?=E6=9B=B4=E6=96=B0=E7=AD=96=E7=95=A5?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- train.py => train_mobilenetv2.py | 49 +++++++++++++++++++++++++------- 1 file changed, 38 insertions(+), 11 deletions(-) rename train.py => train_mobilenetv2.py (68%) diff --git a/train.py b/train_mobilenetv2.py similarity index 68% rename from train.py rename to train_mobilenetv2.py index 02f6847..2ae65e3 100644 --- a/train.py +++ b/train_mobilenetv2.py @@ -27,21 +27,20 @@ if __name__=="__main__": # 数据加载 train_dir = './train_data/1/train' - test_dir = 'train_data/1/val' + test_dir = 'train_data/1/test' # 训练轮数 - num_epochs = 10 + num_epochs = 20 - # + #加载数据集 train_dataset = datasets.ImageFolder(train_dir, transform=transform) train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=64, shuffle=True) - - # test_dataset = datasets.ImageFolder(test_dir, transform=transform) test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=64, shuffle=False) # 构建MobileNetV2模型 model = models.mobilenet_v2(pretrained=True) + num_classes = len(train_dataset.classes) model.classifier[1] = nn.Linear(in_features=1280, out_features=num_classes) @@ -50,7 +49,10 @@ if __name__=="__main__": # 定义损失函数和优化器 criterion = nn.CrossEntropyLoss() - optimizer = optim.Adam(model.parameters(), lr=0.00008) + optimizer = optim.Adam(model.parameters(), lr=0.0001) + + # 添加ReduceLROnPlateau调度器 + scheduler = ReduceLROnPlateau(optimizer, mode='min', factor=0.1, patience=2, verbose=False) print(f"train data:{len(train_loader)}") print(f"test data:{len(test_loader)}") @@ -59,10 +61,14 @@ if __name__=="__main__": # 训练模型 print("start training") - temp = time.time() + min_loss = 10000.0 + max_accuracy = 0 + #temp = time.time() for epoch in range(num_epochs): train_start_time = time.time() print(f"turn {epoch + 1}:") + current_lr = optimizer.param_groups[0]['lr'] + print(f"Current learning rate: {current_lr}") model.train() running_loss = 0.0 for inputs, labels in tqdm(train_loader, desc="training", unit="item", ncols=100): @@ -77,29 +83,50 @@ if __name__=="__main__": running_loss += loss.item() train_end_time = time.time() - print(f"Epoch {epoch + 1}/{num_epochs}, Loss: {running_loss / len(train_loader)} train cost:{train_end_time -train_start_time}") + print(f"Epoch {epoch + 1}/{num_epochs} Train loss: {running_loss / len(train_loader)} Train cost:{train_end_time -train_start_time}") # 在测试集上评估模型 test_start_time = time.time() model.eval() correct = 0 total = 0 + val_loss = 0.0 with torch.no_grad(): for inputs, labels in test_loader: inputs, labels = inputs.to(device), labels.to(device) outputs = model(inputs) + loss = criterion(outputs, labels) + val_loss += loss.item() + _, predicted = torch.max(outputs.data, 1) total += labels.size(0) correct += (predicted == labels).sum().item() + + val_loss /= len(test_loader) + + # 更新学习率 + scheduler.step(val_loss) + test_end_time = time.time() accuracy = 100 * correct / total - print(f"Test Accuracy: {accuracy:.2f}% test cost:{ test_end_time - test_start_time }") + print(f"Test Accuracy: {accuracy:.2f}% Test loss:{ val_loss } test cost:{ test_end_time - test_start_time }") # 保存模型 - torch.save(model.state_dict(), f"./model/1/epochs{epoch + 1} {accuracy:.2f}.pt") - print("model saved success") + if val_loss < min_loss or max_accuracy < accuracy: + min_loss = val_loss + max_accuracy = accuracy + torch.save(model.state_dict(), f"./model/1/epochs{epoch + 1} {accuracy:.2f}.pt") + print("model saved") print("all finish") + + ''' + torch.save(model.state_dict(), f"./model/1/epochs{epoch + 1} {accuracy:.2f}.pt") + print("final model save") + ''' + + + print(f"time:{time.time()-start_time}")