diff --git a/.gitignore b/.gitignore index 89d0310..ddf2a61 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,8 @@ +__pycache__/ +*.py[cod] +*.pyo +*.pyd + *.exe *.exe~ *.dll @@ -8,14 +13,19 @@ *.log *.txt -tmp/ -.claude/ +venv/ +env/ +.venv/ + .vscode/ +.idea/ +.claude/ + +nn/dataset/toy/* +nn/dataset/benchmark/* +nn/models/*.pth +nn/models/*.pt build/bin frontend/node_modules frontend/dist - -nn/dataset/* - - diff --git a/nn/__init__.py b/nn/__init__.py new file mode 100644 index 0000000..2453740 --- /dev/null +++ b/nn/__init__.py @@ -0,0 +1 @@ +# Neural Network Package diff --git a/nn/__pycache__/const.cpython-310.pyc b/nn/__pycache__/const.cpython-310.pyc index 901b0ee..47ddbdc 100644 Binary files a/nn/__pycache__/const.cpython-310.pyc and b/nn/__pycache__/const.cpython-310.pyc differ diff --git a/nn/__pycache__/dataloader.cpython-310.pyc b/nn/__pycache__/dataloader.cpython-310.pyc index cc558e1..3299830 100644 Binary files a/nn/__pycache__/dataloader.cpython-310.pyc and b/nn/__pycache__/dataloader.cpython-310.pyc differ diff --git a/nn/__pycache__/dataloader.cpython-313.pyc b/nn/__pycache__/dataloader.cpython-313.pyc index 2074ebd..c55cbec 100644 Binary files a/nn/__pycache__/dataloader.cpython-313.pyc and b/nn/__pycache__/dataloader.cpython-313.pyc differ diff --git a/nn/const.py b/nn/const.py deleted file mode 100644 index c1b0336..0000000 --- a/nn/const.py +++ /dev/null @@ -1,81 +0,0 @@ -#训练多少轮 -epoch = 50 - -#学习率 -lr = 0.0005 - -#一轮多少张图片 -batch_size = 64 - -#训练图片输入尺寸 -input_size = 128 - -#分类 -label_name = [ - "abyssinian", - "cyprus", - "lykoi", - "donskoy", - "chausie", - "european_shorthair", - "turkish_van", - "pixie_bob", - "ragdoll", - "german_rex", - "american_shorthair", - "sokoke", - "khao_manee", - "thai", - "cymric", - "oriental_shorthair", - "cornish_rex", - "burmese", - "savannah", - "american_wirehair", - "peterbald", - "karelian_bobtail", - "tonkinese", - "balinese", - "japanese_bobtail", - "nebelung", - "selkirk_rex", - "persian", - "manx", - "himalayan", - "munchkin", - "bengal", - "turkish_angora", - "vankedisi", - "scottish_fold", - "egyptian_mau", - "ocicat", - "ragamuffin", - "serengeti", - "british_shorthair", - "toyger", - "siberian", - "havana_brown", - "exotic_shorthair", - "bombay", - "korat", - "safari", - "american_bobtail", - "mekong_bobtail", - "korean_bobtail", - "siamese", - "somali", - "devon_rex", - "american_curl", - "ural_rex", - "singapura", - "ukrainian_levkoy", - "maine_coon", - "birman", - "oregon_rex", - "kurilian_bobtail", - "laperm", - "sphynx", - "chartreux", - "russian_blue", - "norwegian_forest_cat", -] \ No newline at end of file diff --git a/nn/const/__init__.py b/nn/const/__init__.py new file mode 100644 index 0000000..3d09d5f --- /dev/null +++ b/nn/const/__init__.py @@ -0,0 +1,44 @@ +mode = "toy" + +""" +epoch 训练多少轮 +lr 学习率 +batch_size 一轮跑多少张图片 +input_size 训练图片输入尺寸 +label_name 分类 +""" + +if mode == "toy": + epoch = 15 + lr = 2e-4 + batch_size = 2 + input_size = 224 + + # 分类 + label_name = [ + "american_shorthair", + "bengal", + "british_shorthair", + "exotic_shorthair", + "maine_coon", + "ragdoll", + "sphynx", + ] + num_classes = len(label_name) +elif mode == "benchmark": + epoch = 15 + lr = 2e-4 + batch_size = 2 + input_size = 224 + + # 分类 + label_name = [ + "american_shorthair", + "bengal", + "british_shorthair", + "exotic_shorthair", + "maine_coon", + "ragdoll", + "sphynx", + ] + num_classes = len(label_name) \ No newline at end of file diff --git a/nn/dataloader/__init__.py b/nn/dataloader/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/nn/dataloader.py b/nn/dataloader/dataloader.py similarity index 78% rename from nn/dataloader.py rename to nn/dataloader/dataloader.py index 4422a99..1a3f371 100644 --- a/nn/dataloader.py +++ b/nn/dataloader/dataloader.py @@ -3,7 +3,7 @@ import glob from torchvision import transforms from torch.utils.data import DataLoader, Dataset from PIL import Image -from const import label_name, input_size, batch_size +from nn.const import label_name, input_size, batch_size label_dict = {} @@ -13,21 +13,13 @@ for idx, name in enumerate(label_name): def default_loader(path): - img = Image.open(path).convert('RGB') - w, h = img.size - - if w > 200: # 如果宽度超过200，可能是错误数据，等比缩放到162 - ratio = 162 / w - new_h = int(h * ratio) - img = img.resize((162, new_h), Image.BILINEAR) - - return img + return Image.open(path).convert('RGB') train_transform = transforms.Compose([ transforms.Resize((input_size)), transforms.CenterCrop(input_size), - transforms.RandomHorizontalFlip(p=0.5), # 50% 的概率（p=0.5）水平翻转图片 + transforms.RandomHorizontalFlip(p=0.5), # 50%的概率（p=0.5）水平翻转图片 transforms.RandomRotation(10), # 轻微旋转 transforms.ColorJitter(brightness=0.1, contrast=0.1), transforms.ToTensor(), @@ -69,8 +61,9 @@ class MyDataset(Dataset): return len(self.imgs) -im_train_list = glob.glob("dataset/train/*/*.jpg") -im_test_list = glob.glob("dataset/test/*/*.jpg") +dataset_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +im_train_list = glob.glob(os.path.join(dataset_root, "dataset", "toy", "train", "*", "*.jpg")) +im_test_list = glob.glob(os.path.join(dataset_root, "dataset", "toy", "test", "*", "*.jpg")) train_dataset = MyDataset(im_train_list, transform=train_transform) diff --git a/nn/model/resnet_epoch_31.pth b/nn/model/resnet_epoch_31.pth deleted file mode 100644 index 34bcdb9..0000000 Binary files a/nn/model/resnet_epoch_31.pth and /dev/null differ diff --git a/nn/resnet.py b/nn/nets/resnet.py similarity index 97% rename from nn/resnet.py rename to nn/nets/resnet.py index d16e0ea..71ad3b6 100644 --- a/nn/resnet.py +++ b/nn/nets/resnet.py @@ -1,5 +1,6 @@ import torch.nn as nn import torch.nn.functional as F +from nn.const import num_classes class ResBlock(nn.Module): @@ -44,7 +45,7 @@ class ResNet(nn.Module): return nn.Sequential(*layer_list) - def __init__(self, num_classes=67): + def __init__(self): super(ResNet, self).__init__() self.in_channel = 32 diff --git a/nn/nets/resnet18.py b/nn/nets/resnet18.py new file mode 100644 index 0000000..429c8ce --- /dev/null +++ b/nn/nets/resnet18.py @@ -0,0 +1,19 @@ +import torch.nn as nn +from torchvision import models +from nn.const import num_classes + + +class resnet18(nn.Module): + def __init__(self): + super(resnet18, self).__init__() + self.model = models.resnet18(weights='IMAGENET1K_V1') + self.num_features = self.model.fc.in_features + self.model.fc = nn.Linear(self.num_features, num_classes) + + def forward(self, x): + out = self.model(x) + return out + + +def pytorch_resnet18(): + return resnet18() \ No newline at end of file diff --git a/nn/test.py b/nn/test/test_resnet.py similarity index 72% rename from nn/test.py rename to nn/test/test_resnet.py index e00d1ce..5d70a80 100644 --- a/nn/test.py +++ b/nn/test/test_resnet.py @@ -4,17 +4,16 @@ import torch from torchvision import transforms from PIL import Image import numpy as np -# from resnet import resnet -from torchvision.models import resnet18 -from const import label_name, input_size +from nn.nets.resnet import resnet +from nn.const import label_name, input_size def test(): device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print(device) - net = resnet18(weights=None) - net.load_state_dict(torch.load("./model/resnet_epoch_31.pth", weights_only=True)) + net = resnet() + net.load_state_dict(torch.load("./model/resnet_epoch_15.pth", weights_only=True)) im_list = glob.glob("./dataset/test/*/*.jpg") np.random.shuffle(im_list) @@ -22,8 +21,10 @@ def test(): net.to(device) test_transform = transforms.Compose([ - transforms.Resize((input_size, input_size)), - transforms.ToTensor() + 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: @@ -35,7 +36,6 @@ 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) diff --git a/nn/test/test_resnet18.py b/nn/test/test_resnet18.py new file mode 100644 index 0000000..d4890fe --- /dev/null +++ b/nn/test/test_resnet18.py @@ -0,0 +1,55 @@ +import cv2 +import glob +import torch +from torchvision import transforms +from PIL import Image +import numpy as np +from nn.nets.resnet18 import resnet18 +from nn.const import label_name, input_size + + +def test(): + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + print(device) + + net = resnet18() + net.load_state_dict(torch.load("./model/resnet_epoch_14.pth", weights_only=True)) + + im_list = glob.glob("./dataset/test/*/*.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) + # 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) + cv2.imshow("img", img) + cv2.waitKey(0) + + +if __name__ == "__main__": + test() diff --git a/nn/train.py b/nn/train/train_resnet.py similarity index 83% rename from nn/train.py rename to nn/train/train_resnet.py index 2081c89..06cc961 100644 --- a/nn/train.py +++ b/nn/train/train_resnet.py @@ -1,23 +1,21 @@ import os import torch -# from resnet import resnet -from torchvision.models import resnet18 -from dataloader import train_dataloader -from const import epoch, lr, batch_size +from nn.nets.resnet import resnet +from nn.dataloader.dataloader import train_dataloader +from nn.const import epoch, lr, batch_size def train(): device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print("device: ", device) - # net = resnet().to(device) - net = resnet18(weights=None).to(device) + net = resnet().to(device) loss_func = torch.nn.CrossEntropyLoss() optimizer = torch.optim.Adam(net.parameters(), lr=lr) - scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.5) + scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.5) for e in range(epoch): print("epoch: ", e) diff --git a/nn/train/train_resnet18.py b/nn/train/train_resnet18.py new file mode 100644 index 0000000..69acb4f --- /dev/null +++ b/nn/train/train_resnet18.py @@ -0,0 +1,51 @@ +import os +import torch +from nn.nets.resnet18 import resnet18 +from nn.dataloader.dataloader import train_dataloader +from nn.const import epoch, lr, batch_size + + +def train(): + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + print("device: ", device) + + net = resnet18().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']) + + if not os.path.exists("./models"): + os.makedirs("./models") + + torch.save(net.state_dict(), "./models/resnet18_epoch_{}.pth".format(e + 1)) + + +if __name__ == "__main__": + train() \ No newline at end of file