From 4d0b67831f402d3b925557424044e11d01c93dd5 Mon Sep 17 00:00:00 2001
From: Lv Yan <2839719742@qq.com>
Date: Fri, 30 Jul 2021 22:13:06 +0800
Subject: [PATCH] add the codes of train and test

---
 test.py  |  53 ++++++++++++++++++++++++
 train.py | 120 +++++++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 173 insertions(+)
 create mode 100644 test.py
 create mode 100644 train.py

diff --git a/test.py b/test.py
new file mode 100644
index 0000000..06b793a
--- /dev/null
+++ b/test.py
@@ -0,0 +1,53 @@
+import torch
+import torchvision
+import torchvision.transforms as transforms
+from torch.utils.data import DataLoader
+
+from nets.ghostnet import ghostnet
+
+transform_test = transforms.Compose([
+    transforms.CenterCrop(960),
+    transforms.Resize(512),
+    transforms.ToTensor()
+])
+
+testset = torchvision.datasets.ImageFolder(root='./dataset/test', transform=transform_test)
+testloader = DataLoader(testset, batch_size=8, shuffle=False, num_workers=8)
+
+# 使用gpu
+use_cuda = torch.cuda.is_available()
+device = torch.device("cuda" if use_cuda else "cpu")
+
+# 模型定义
+net = ghostnet().to(device)
+
+
+def test():
+    print("Load checkpoint...")
+    net.load_state_dict(torch.load('./model_MESSIDOR/model_MESSIDOR_3000.pkl'))
+
+    print("test begin！")
+    with torch.no_grad():
+        accuracy = 0
+        total = 0
+        for data in testloader:
+            # 开始测试
+            net.eval()
+
+            images, labels = data
+            images, labels = images.to(device), labels.to(device)
+
+            outputs = net(images)
+
+            _, predicted = torch.max(outputs.data, 1)  # 返回每一行中最大值的那个元素，且返回其索引(得分高的那一类)
+            total += labels.size(0)
+
+            accuracy += (predicted == labels).sum()
+        accuracy = accuracy.item()  # acc为tensor, total为int型, acc/total自动将小数点后数字归零. 所以需要先将acc通过.item()转化成int型.
+
+        # 输出测试准确率
+        print('测试准确率为: {:.2f}%'.format(100 * accuracy / total))
+
+
+if __name__ == '__main__':
+    test()
\ No newline at end of file
diff --git a/train.py b/train.py
new file mode 100644
index 0000000..f107176
--- /dev/null
+++ b/train.py
@@ -0,0 +1,120 @@
+import torch
+import torchvision
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import DataLoader
+import torchvision.transforms as transforms
+import os
+import sys
+
+from nets.ghostnet import ghostnet
+
+sys.path.append(os.getcwd())
+os.environ["CUDA_VISIBLE_DEVICES"] = "0"  # 指定一块gpu为可见
+
+# #############创建数据加载器###################
+print('data loaded begin!')
+# 预处理
+data_transform = transforms.Compose([
+    transforms.CenterCrop(960),
+    transforms.Resize(512),
+    transforms.ToTensor()
+])
+# 加载训练集数据
+train_dataset = torchvision.datasets.ImageFolder(root='./datasets/train', transform=data_transform)
+train_data = DataLoader(train_dataset, batch_size=8, shuffle=True, num_workers=8, drop_last=True)
+# 加载验证集数据
+val_dataset = torchvision.datasets.ImageFolder(root='./datasets/val', transform=data_transform)
+val_data = DataLoader(val_dataset, batch_size=8, shuffle=True, num_workers=8, drop_last=True)
+
+print(type(train_data))
+print('data loaded done!')
+
+# ##################创建网络模型###################
+model = ghostnet()
+
+# ##############训练#################
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+if torch.cuda.device_count() > 1:
+    model = nn.DataParallel(model)
+model = model.to(device)
+criterion = nn.CrossEntropyLoss()
+optimizer = optim.Adam(model.parameters(), weight_decay=0.001)
+
+nums_epoch = 3001  # 训练epoch的数量，动态调整
+
+print('training begin!')
+# 开始训练
+train_losses = []
+train_acces = []
+val_losses = []
+val_acces = []
+S = []
+
+for epoch in range(nums_epoch):
+    train_loss = 0
+    train_acc = 0
+    model = model.train()
+    # 训练开始
+    print('Epoch ' + str(epoch + 1) + ' begin!')
+    for img_train, label_train in train_data:
+        img_train = img_train.to(device)
+        label_train = label_train.to(device)
+
+        # 前向传播
+        out = model(img_train)
+        optimizer.zero_grad()
+        loss = criterion(out, label_train)
+        # 反向传播
+        loss.backward()
+
+        optimizer.step()
+
+        # 记录误差
+        train_loss += loss.item()
+        # 计算分类的准确率
+        _, pred = out.max(1)
+        num_correct = (pred == label_train).sum().item()
+        acc = num_correct / img_train.shape[0]
+        train_acc += acc
+    train_losses.append(train_loss / len(train_data))
+    train_acces.append(train_acc / len(train_data))
+    print('Epoch' + str(epoch + 1) + ' Train  done!')
+
+    val_loss = 0
+    val_acc = 0
+    # 验证开始
+    print('Epoch' + str(epoch + 1) + ' Val  begin!')
+    for img_val, label_val in val_data:
+        img_val = img_val.to(device)
+        label_val = label_val.to(device)
+
+        out = model(img_val)
+        loss = criterion(out, label_val)
+
+        # 记录误差
+        val_loss += loss.item()
+        # 计算分类的准确率
+        _, pred = out.max(1)
+        num_correct = (pred == label_val).sum().item()
+        acc = num_correct / img_val.shape[0]
+        val_acc += acc
+    val_losses.append(val_loss / len(val_data))
+    val_acces.append(val_acc / len(val_data))
+    print('Epoch' + str(epoch + 1) + ' Val  done!')
+
+    print('Epoch {} ,Train Loss: {} ,Train  Accuracy: {} ,Test Loss: {} ,Test Accuracy: {}'.format(
+        epoch + 1, train_loss / len(train_data), train_acc / len(train_data), val_loss / len(val_data),
+        val_acc / len(val_data)))
+
+    s = ('Epoch {} ,Train Loss: {} ,Train  Accuracy: {} ,Test Loss: {} ,Test Accuracy: {}'.format(
+        epoch + 1, train_loss / len(train_data), train_acc / len(train_data), val_loss / len(val_data),
+        val_acc / len(val_data)))
+
+    S.append(s)
+
+    if epoch % 100 == 0:
+        torch.save(model.state_dict(), './models/model_MESSIDOR/model_MESSIDOR_{}.pkl'.format(epoch))
+
+    print('model saved done!')
+
-- 
GitLab