From 17b5459ce151cf3e9c5d8bac0129975392716877 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=A4=8F=E4=B8=9C=E4=BA=AE?= Date: Wed, 5 Aug 2026 16:59:16 +0800 Subject: [PATCH] convnext --- core/const/const.py | 2 +- core/nets/convnext_tiny.py | 19 ++++++++++ core/test/test_convnext_tiny.py | 59 +++++++++++++++++++++++++++++++ core/test/test_resnet18.py | 7 ---- core/train/train_convnext_tiny.py | 53 +++++++++++++++++++++++++++ 5 files changed, 132 insertions(+), 8 deletions(-) create mode 100644 core/nets/convnext_tiny.py create mode 100644 core/test/test_convnext_tiny.py create mode 100644 core/train/train_convnext_tiny.py diff --git a/core/const/const.py b/core/const/const.py index 3399357..eeab665 100644 --- a/core/const/const.py +++ b/core/const/const.py @@ -26,7 +26,7 @@ if mode == "toy": ] num_classes = len(label_name) elif mode == "benchmark": - epoch = 50 + epoch = 100 # convnext->100, resnet18->50 lr = 1e-4 batch_size = 8 input_size = 224 diff --git a/core/nets/convnext_tiny.py b/core/nets/convnext_tiny.py new file mode 100644 index 0000000..f8d0839 --- /dev/null +++ b/core/nets/convnext_tiny.py @@ -0,0 +1,19 @@ +import torch.nn as nn +from torchvision import models +from core.const.const import num_classes + + +class ConvNeXtTiny(nn.Module): + def __init__(self): + super(ConvNeXtTiny, self).__init__() + self.model = models.convnext_tiny(weights='IMAGENET1K_V1') + self.num_features = self.model.classifier[2].in_features + self.model.classifier[2] = nn.Linear(self.num_features, num_classes) + + def forward(self, x): + out = self.model(x) + return out + + +def pytorch_convnext_tiny(): + return ConvNeXtTiny() diff --git a/core/test/test_convnext_tiny.py b/core/test/test_convnext_tiny.py new file mode 100644 index 0000000..23681db --- /dev/null +++ b/core/test/test_convnext_tiny.py @@ -0,0 +1,59 @@ +import os +import cv2 +import glob +import torch +from torchvision import transforms +from PIL import Image +import numpy as np +from core.nets.convnext_tiny import pytorch_convnext_tiny +from core.const import mode, label_name, input_size + + +def test(): + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + print(device) + + script_dir = os.path.dirname(os.path.abspath(__file__)) + core_dir = os.path.dirname(script_dir) + model_dir = os.path.join(core_dir, "models") + dataset_dir = os.path.join(core_dir, "dataset", mode, "test") + + print("model_dir", model_dir) + + net = pytorch_convnext_tiny() + net.load_state_dict(torch.load(os.path.join(model_dir, "convnext_tiny_epoch_100.pth"), weights_only=True)) + + im_list = glob.glob(os.path.join(dataset_dir, "*", "*.jpg")) + np.random.shuffle(im_list) + + net.to(device) + + test_transform = transforms.Compose([ + transforms.Resize(input_size), + transforms.CenterCrop(input_size), + transforms.ToTensor(), + transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) + ]) + + for im_path in im_list: + net.eval() + im_data = Image.open(im_path) + + inputs = test_transform(im_data) + inputs = torch.unsqueeze(inputs, dim=0) + + inputs = inputs.to(device) + outputs = net.forward(inputs) + + _, pred = torch.max(outputs.data, dim=1) + print(label_name[pred.cpu().numpy()[0]], " ", im_path) + + img = np.asarray(im_data) + img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR) + img = cv2.resize(img, (200, 200)) + cv2.imshow("img", img) + cv2.waitKey(0) + + +if __name__ == "__main__": + test() diff --git a/core/test/test_resnet18.py b/core/test/test_resnet18.py index 457e16e..c994300 100644 --- a/core/test/test_resnet18.py +++ b/core/test/test_resnet18.py @@ -21,10 +21,8 @@ def test(): print("model_dir", model_dir) net = resnet18() - # net.load_state_dict(torch.load("./models/resnet18_epoch_100.pth", weights_only=True)) net.load_state_dict(torch.load(os.path.join(model_dir, "resnet18_epoch_50_bak2.pth"), weights_only=True)) - # im_list = glob.glob("./dataset/test/*/*.jpg") im_list = glob.glob(os.path.join(dataset_dir, "*", "*.jpg")) np.random.shuffle(im_list) @@ -46,15 +44,10 @@ def test(): inputs = inputs.to(device) outputs = net.forward(inputs) - # print("outputs", outputs) _, pred = torch.max(outputs.data, dim=1) print(label_name[pred.cpu().numpy()[0]], " ", im_path) - prob, pred = torch.topk(outputs.data, k=3, dim=1) - # for i in range(3): - # print(label_name[pred[0, i].item()], " ", prob[0, i].item(), " ", im_path) - img = np.asarray(im_data) img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR) img = cv2.resize(img, (200, 200)) diff --git a/core/train/train_convnext_tiny.py b/core/train/train_convnext_tiny.py new file mode 100644 index 0000000..1f5344f --- /dev/null +++ b/core/train/train_convnext_tiny.py @@ -0,0 +1,53 @@ +import os +import torch +from core.nets.convnext_tiny import pytorch_convnext_tiny +from core.dataloader.dataloader import train_dataloader +from core.const import epoch, lr, batch_size + + +def train(): + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + print("device: ", device) + + net = pytorch_convnext_tiny().to(device) + + loss_func = torch.nn.CrossEntropyLoss() + + optimizer = torch.optim.Adam(net.parameters(), lr=lr) + + scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=20, eta_min=1e-6) + + for e in range(epoch): + print("epoch: ", e) + net.train() + + for i, data in enumerate(train_dataloader): + inputs, labels = data + inputs, labels = inputs.to(device), labels.to(device) + + outputs = net(inputs) + + loss = loss_func(outputs, labels) + + optimizer.zero_grad() + loss.backward() + optimizer.step() + + _, pred = torch.max(outputs, dim=1) + correct = pred.eq(labels.data).cpu().sum() + + print("step: ", i, "loss: ", loss.item(), "correct: ", 1.0 * correct / batch_size) + + scheduler.step() + print("lr: ", optimizer.state_dict()['param_groups'][0]['lr']) + + script_dir = os.path.dirname(os.path.abspath(__file__)) + model_dir = os.path.join(script_dir, "..", "models") + if not os.path.exists(model_dir): + os.makedirs(model_dir) + + torch.save(net.state_dict(), os.path.join(model_dir, "convnext_tiny_epoch_{}.pth".format(e + 1))) + + +if __name__ == "__main__": + train()