xiadongliang/SortMeow
1
1
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

工程化

Browse Source
This commit is contained in:
夏东亮
2026-07-09 17:48:39 +08:00
parent 3d2744c5b1
commit 39a25e488e
16 changed files with 207 additions and 116 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+16 -6
View File
@@ -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/*
nn/__init__.py
+1
View File
@@ -0,0 +1 @@
# Neural Network Package
nn/__pycache__/const.cpython-310.pyc
BIN
View File
Binary file not shown.
nn/__pycache__/dataloader.cpython-310.pyc
BIN
View File
Binary file not shown.
nn/__pycache__/dataloader.cpython-313.pyc
BIN
View File
Binary file not shown.
nn/const.py
-81
View File
@@ -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",
]
nn/const/__init__.py
+44
View File
@@ -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)
nn/dataloader/__init__.py
View File
nn/dataloader.py → nn/dataloader/dataloader.py
+6 -13
View File
@@ -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)
nn/model/resnet_epoch_31.pth
BIN
View File
Binary file not shown.
nn/resnet.py → nn/nets/resnet.py
+2 -1
View File
@@ -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
nn/nets/resnet18.py
+19
View File
@@ -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()
nn/test.py → nn/test/test_resnet.py
+8 -8
View File
@@ -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)
nn/test/test_resnet18.py
+55
View File
@@ -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()
nn/train.py → nn/train/train_resnet.py
+5 -7
View File
@@ -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)
nn/train/train_resnet18.py
+51
View File
@@ -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()