convnext

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

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()

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()

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))

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()