0基础肝ai比赛用的Pytorch框架（codex写的代码）

Pytorch框架（图像）

文件结构

NS-0
├─ data
│  ├─ train.csv
│  ├─ test.csv
│  └─ images
├─ outputs
│  └─ model.pth
└─ src
   ├─ train.py
   └─ predict.py

train.py

一、导入工具

普通 Python 工具

import os
import json
import random
import pandas as pd
from PIL import Image

•os：处理文件路径、创建文件夹。
•json：处理 JSON 数据，这里其实没用到。
•random：设置随机数，保证结果尽量稳定。
•pandas as pd：读取和保存 CSV 表格。
•Image：打开图片文件。

PyTorch 相关工具

import torch
from torch import nn
from torch.utils.data import Dataset, DataLoader
from torchvision import models, transforms

•torch：深度学习框架本体。
•nn：神经网络模块，比如损失函数、全连接层。
•Dataset：定义“怎么读取一条数据”。
•DataLoader：把数据一批一批送给模型。
•models：现成的模型，比如 ResNet18。
•transforms：图片预处理，比如缩放、转成张量、归一化

二、配置参数

BASE_DIR = os.path.dirname(os.path.dirname(__file__)) #基目录

TRAIN_CSV = os.path.join(BASE_DIR, "data", "train.csv") #训练标签表在 data/train.csv
TRAIN_IMG_DIR = os.path.join(BASE_DIR, "data", "train_images") #训练图片在 data/train_images（视csv词条决定）
OUTPUT_DIR = os.path.join(BASE_DIR, "outputs") #输出目录
OUTPUT_MODEL = os.path.join(OUTPUT_DIR, "model.pth") #模型保存在outputs/model.pth
LABEL_MAP_PATH = os.path.join(OUTPUT_DIR, "id_to_label.json") #标签映射文件也要保存

IMAGE_COL = "image" #看CSV里，图片文件名那一列叫 image
LABEL_COL = "label" #标签那一列叫 label

三、训练参数

NUM_CLASSES = None #一共有几个类别（建议让程序自己算）
BATCH_SIZE = 4 #每次拿几张图片训练
EPOCHS = 1 #训练几轮
LR = 1e-4 #学习率，控制模型每次更新幅度

四、class ImageDataset(Dataset):

继承Dataset，告诉 PyTorch “我的图片在哪里，标签在哪里，怎么读取一张图片”。

初始化

    def __init__(self, csv_path, image_dir, transform=None):
        self.df = pd.read_csv(csv_path)
        self.image_dir = image_dir
        self.transform = transform
		#如果是训练集，就把标签转成数字
        labels = sorted(self.df[LABEL_COL].unique())
        self.label_to_id = {label: i for i, label in enumerate(labels)}

如果是测试集，没有真实标签，就不用做标签转换。

返回数据集长度

def __len__(self):
	return len(self.df)

读取一张图片

def __getitem__(self, index):
	row = self.df.iloc[index]

	image_path = os.path.join(self.image_dir, row[IMAGE_COL])
	image = Image.open(image_path).convert("RGB")

	if self.transform is not None:
		image = self.transform(image)

	label = self.label_to_id[row[LABEL_COL]]
	return image, label

五、main函数

作用：串联所有过程

选择 CPU/GPU
准备图片处理
创建 Dataset
创建 DataLoader
创建模型
创建 loss 和 optimizer
开始训练
保存模型

1.训练设备

device = "cuda" if torch.cuda.is_available() else "cpu"
    print("Using device:", device)

2.图片预处理

train_transform = transforms.Compose([
    transforms.Resize((224, 224)), #统一缩放到 224x224
    transforms.RandomRotation(5), #仅训练：随机旋转
    transforms.ColorJitter(brightness=0.2, contrast=0.2), #仅训练：随机改变亮度和对比度
    transforms.ToTensor(), #把图片转成 PyTorch 能处理的数字格式
    transforms.Normalize(  #标准化图片数值，让模型更容易训练
        mean=[0.485, 0.456, 0.406],
        std=[0.229, 0.224, 0.225]
    ),
])

3.创建训练数据集和加载器

dataset = ImageDataset(
        csv_path=TRAIN_CSV,
        image_dir=TRAIN_IMG_DIR,
        transform=transform,
    )
    
dataloader = DataLoader(
        dataset,
        batch_size=BATCH_SIZE, #每次拿BATCH_SIZE张图片
        shuffle=True, #每轮训练前打乱顺序
    )

程序自己算num_classes

num_classes = len(dataset.label_to_id)
#print("Classes:", dataset.label_to_id)
#print("Num classes:", num_classes)

4.创建模型

main:

model = build_model(num_classes).to(device)

def build_model(num_classes):
    #Create an image classification model.
    model = models.resnet18(weights=None)
    model.fc = nn.Linear(model.fc.in_features, num_classes)
    return model

# model = models.resnet18(weights=models.ResNet18_Weights.IMAGENET1K_V1) #加载一个已经在 ImageNet 上训练过的ResNet18 模型
model = model.to(device) #把模型放到 GPU 或 CPU 上

损失函数和优化器

loss_fn = nn.CrossEntropyLoss() #损失函数：模型预测得离正确答案有多远
optimizer = torch.optim.Adam(model.parameters(), lr=LR) #优化器：负责根据错误程度调整模型参数，让模型下次预测更准

5.训练模型（核心）

for epoch in range(EPOCHS):

每一轮：

model.train() #切换到训练模式
total_loss = 0 #记录这一轮的总损失

for batch_index, (images, labels) in enumerate(dataloader): #一批一批读取训练图片和标签
    images = images.to(device)
    labels = labels.to(device) #把图片和标签放到 GPU 或 CPU 上

    outputs = model(images) #把图片送进模型，得到预测结果
    loss = loss_fn(outputs, labels) #计算模型预测和真实标签之间的差距

    optimizer.zero_grad() #清空上一次的梯度（每次更新模型前，先把旧的计算痕迹清掉）
    loss.backward() #反向传播（计算模型里面哪些参数导致了错误，以及应该怎么调整）
    optimizer.step() #真正更新模型参数

    total_loss += loss.item() #把这一批的损失记录下来

    if batch_index % 10 == 0: #每10组汇报一次进度
        print(f"Batch {batch_index}/{len(dataloader)}, loss={loss.item():.4f}") #如果训练正常，loss 通常会慢慢下降

        return total_loss / len(dataloader)

即核心流程是：

拿图片和标签

送进模型

算错误

清空旧梯度

反向传播

更新参数

6.保存模型

id_to_label = {i: label for label, i in dataset.label_to_id.items()} #标签保存部分
os.makedirs(OUTPUT_DIR, exist_ok=True)
torch.save(model.state_dict(), OUTPUT_MODEL)
with open(LABEL_MAP_PATH, "w", encoding="utf-8") as f:  #标签保存部分
	json.dump(id_to_label, f, ensure_ascii=False, indent=2)

predict.py

predict.py的功能：

设置路径
读取 test.csv
创建测试 Dataset
创建同样结构的模型
加载 outputs/model.pth
对测试图片预测
把预测结果保存成 outputs/result.csv

posted @ 2026-05-29 15:44 清露宵偃阅读(6) 评论(0) 收藏举报

刷新页面返回顶部