这只是介绍怎么简单使用，并不介绍densenet。

qq群机器人之前有发送pixiv的二次元图片的功能，但是发送的有的并不是我喜欢的，所以想用ai来分类一下，所以就有了此项目。

安装pytorch

首先保证你有英伟达的卡，然后安装cuda，没有的话能接受cpu那非常慢的训练速度也可以用cpu

cuda 下载

没有英伟达的卡不需要安装cuda

pytorch start

选择对应平台，对应cuda版本 or cpu。

准备代码

建立文件夹

项目\
 --dataset
    -- train   #这是训练集
        -- a.喜欢的
        -- b.不喜欢的
    -- validation  # 这是测试集
        -- a.喜欢的
        -- b.不喜欢的
 main.py

训练集就是用来训练模型的，测试集就是用来检验模型正确率的，测试集不要包含训练集相同的图片。

main.py

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision.transforms as transforms
import torchvision.datasets as datasets
from torch.autograd import Variable
from torchvision.models import densenet169

# 设置全局参数
model_lr = 1e-4
batch_size = 32
epochs = 150
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 图像预处理
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
transform_test = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])

def adjust_learning_rate(optimizer, epoch):
    """Sets the learning rate to the initial lr decayed by 10 every 30 epochs"""
    model_lr_new = model_lr * (0.1 ** (epoch // 50))
    print("lr:", model_lr_new)
    for param_group in optimizer.param_groups:
        param_group['lr'] = model_lr_new
    print("lr changed")

# 定义训练过程
def train(model, device, train_loader, optimizer, epoch):
    model.train()
    sum_loss = 0
    total_num = len(train_loader.dataset)
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = Variable(data).to(device), Variable(target).to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()
        sum_loss += loss.item() * data.size(0)
    print('Train Epoch: {} \tLoss: {:.6f}'.format(
        epoch, sum_loss / total_num))

# 定义验证过程
def validate(model, device, test_loader):
    model.eval()
    val_loss = 0
    correct = 0
    total = 0
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            loss = criterion(output, target)
            val_loss += loss.item() * data.size(0)
            _, predicted = torch.max(output.data, 1)
            total += target.size(0)
            correct += (predicted == target).sum().item()
    val_loss /= len(test_loader.dataset)
    val_accuracy = 100 * correct / total
    print('Validation Loss: {:.4f}, Validation Accuracy: {:.2f}%'.format(
        val_loss, val_accuracy))

if __name__ == '__main__':
    import torch.multiprocessing
    torch.multiprocessing.freeze_support()

    # 读取数据
    dataset_train = datasets.ImageFolder('./dataset/train', transform)
    dataset_test = datasets.ImageFolder('./dataset/validation', transform_test)
    train_loader = torch.utils.data.DataLoader(dataset_train, batch_size=batch_size, shuffle=True, num_workers=8, pin_memory=True)
    test_loader = torch.utils.data.DataLoader(dataset_test, batch_size=batch_size, shuffle=False, num_workers=8, pin_memory=True)

    # 设置模型
    criterion = nn.CrossEntropyLoss()
    model_ft = densenet169(pretrained=True)
    num_ftrs = model_ft.classifier.in_features
    model_ft.classifier = nn.Linear(num_ftrs, 2)
    model_ft.to(device)
    optimizer = optim.Adam(model_ft.parameters(), lr=model_lr)

    # 训练和验证循环
    for epoch in range(1, epochs + 1):
        adjust_learning_rate(optimizer, epoch)
        train(model_ft, device, train_loader, optimizer, epoch)
        if epoch % 10 == 0:
            validate(model_ft, device, test_loader)

    torch.save(model_ft, 'densenet_model.pth')

一共训练150轮，训练完之后保存模型到当前文件夹的 densenet_model.pth，运行main.py就可以开始

使用

我简单用flask写了一个接口。

import io
from flask import Flask, request
import torch
import torchvision.transforms as transforms
from torchvision.models import densenet121
import torch.nn as nn
from PIL import Image
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])

# 加载模型结构
model = torch.load('./densenet_model.pth', weights_only=False)

model.eval()  # 设置为评估模式，关闭一些如Dropout等训练时才需要的操作

app = Flask(__name__)
@app.route('/')
def hello_world():
    return 'Hello World!'

@app.route('/image', methods=['POST'])
def image():
    img = request.files.get("image")
    if img:
        file_bytes = io.BytesIO(img.read())
        image = Image.open(file_bytes).convert('RGB')  # 打开图片并转换为RGB格式（如果需要）
        image = transform(image).unsqueeze(0)  # 应用预处理，并增加一个批次维度（因为模型输入要求有批次维度）
        image = image.to(torch.device('cuda' if torch.cuda.is_available() else 'cpu'))  # 将图片数据移动到相应设备（GPU或CPU）
        output = model(image)
        _, predicted = torch.max(output, 1)  # 获取预测的类别索引
        img.seek(0)
        if predicted.item() == 1:
            print("图片：", img, "属于不喜欢的类别")
            img.save("./upload/b.不喜欢/" + img.filename)

            return "不喜欢"
        else:
            print("图片：", img, "属于喜欢的类别")
            img.save("./upload/a.喜欢/" + img.filename)
            return "喜欢"
    return "请选择图片"

if __name__ == "__main__":
    app.run(port=4005)

保存到api.py，然后运行就好了

这个运行起来会在4005端口上开一个服务，访问 ip:4005/image post方式 image为图片。 会返回是否喜欢。(并未使用json，只是我简单用用)。

这只是一个简单的介绍。其他方面有问题可以评论亦或者联系我。