UNet - 训练数据train

目录

1. train 训练数据

2. Loss 值

3. 完整代码


1. train 训练数据

训练的代码只是在之前图像分类的基础上做了一些更改,具体的可以看下面的文章

pytorch 搭建 LeNet 网络对 CIFAR-10 图片分类https://blog.csdn.net/qq_44886601/article/details/127498256

首先,导入之前定义的UNet 网络

然后,加载训练集和测试集

因为加载数据集被重写过,所以这里传入的是训练的图像,然后根据里面的replace就能找到对应的标签

这里训练的时候可以将数据打乱,测试的时候没有必要,batch_size 因为电脑硬件的问题设置成2,再大的话这里内存就会不够了

 

然后定义优化器和损失函数,这里用的是BCE加上sigmoid的损失函数

训练的时候,要将模式改为train模式,然后训练的步骤很常规

梯度清零->前向传播->计算损失函数->反向传播->更新参数

 

这里测试的时候有些区别

因为这里UNet 网络的输出是一幅图像,而之前将label改为了二值图像(归一化后是0 1)。所以这里计算准确率的时候,将预测的图像也变为二值图像,计算准确率用的是对应图像像素点的灰度值是否相等的方法

 

最后保留最好准确率的那个参数就行了

 

2. Loss 值

这是跑了20 个epoch的输出

 

3. 完整代码

from model import UNet                  # 导入Unet 网络
from dataset import Data_Loader         # 数据处理
from torch import optim
import torch.nn as nn
import torch


# 网络训练模块
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')   # GPU or CPU
print(device)
net = UNet(in_channels=1, num_classes=1)                                # 加载网络
net.to(device)                                                          # 将网络加载到device上

# 加载训练集
train_path = "./data/train/image"
trainset = Data_Loader(train_path)
train_loader = torch.utils.data.DataLoader(dataset=trainset,batch_size=2,shuffle=True)

# len(trainset)  样本总数:21

# 加载测试集
test_path = "./data/test/image"
testset = Data_Loader(test_path)
test_loader = torch.utils.data.DataLoader(dataset=testset,batch_size=2)

optimizer = optim.RMSprop(net.parameters(),lr = 0.000001,weight_decay=1e-8,momentum=0.9)     # 定义优化器
criterion = nn.BCEWithLogitsLoss()                                                           # 定义损失函数

save_path = './UNet.pth'        # 网络参数的保存路径
best_acc = 0.0                  # 保存最好的准确率


for epoch in range(20):

    net.train()     # 训练模式
    running_loss = 0.0

    for image,label in train_loader:                   # 读取数据和label

        optimizer.zero_grad()                          # 梯度清零
        pred = net(image.to(device))                   # 前向传播

        loss = criterion(pred, label.to(device))       # 计算损失
        loss.backward()                                # 反向传播
        optimizer.step()                               # 梯度下降

        running_loss += loss.item()                    # 计算损失和

    net.eval()  # 测试模式
    acc = 0.0   # 正确率
    total = 0
    with torch.no_grad():
        for test_image, test_label in test_loader:

            outputs = net(test_image.to(device))     # 前向传播

            outputs[outputs >= 0] = 1  # 将预测图片转为二值图片
            outputs[outputs < 0] = 0

            acc += (outputs == test_label.to(device)).sum().item() / (480*480)     # 计算预测图片与真实图片像素点一致的精度:acc = 相同的 / 总个数
            total += test_label.size(0)

    accurate = acc / total  # 计算整个test上面的正确率
    print('[epoch %d] train_loss: %.3f  test_accuracy: %.3f %%' %
          (epoch + 1, running_loss, accurate*100))

    if accurate > best_acc:     # 保留最好的精度
        best_acc = accurate
        torch.save(net.state_dict(), save_path)     # 保存网络参数

好的,以下是基于自己的数据训练Unet模型的步骤: 1. 准备数据集:收集和整理训练数据,确保数据集中包含正确的标注信息。可以使用开源数据集,也可以自己制作数据集。 2. 数据增强:使用数据增强技术来扩充数据集,例如旋转、翻转、裁剪等操作,以增加数据集的多样性和数量。 3. 数据预处理:对数据进行预处理,例如对图像进行归一化、缩放等操作。 4. 构建模型:使用Keras或者PyTorch等框架,构建Unet模型,并根据实际需求进行修改和调整。 ```python from tensorflow.keras.models import Model from tensorflow.keras.layers import Input, Conv2D, MaxPooling2D, Dropout, UpSampling2D, concatenate def Unet(input_size): inputs = Input(input_size) # 编码器部分 conv1 = Conv2D(64, 3, activation='relu', padding='same', kernel_initializer='he_normal')(inputs) conv1 = Conv2D(64, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv1) pool1 = MaxPooling2D(pool_size=(2, 2))(conv1) conv2 = Conv2D(128, 3, activation='relu', padding='same', kernel_initializer='he_normal')(pool1) conv2 = Conv2D(128, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv2) pool2 = MaxPooling2D(pool_size=(2, 2))(conv2) conv3 = Conv2D(256, 3, activation='relu', padding='same', kernel_initializer='he_normal')(pool2) conv3 = Conv2D(256, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv3) pool3 = MaxPooling2D(pool_size=(2, 2))(conv3) conv4 = Conv2D(512, 3, activation='relu', padding='same', kernel_initializer='he_normal')(pool3) conv4 = Conv2D(512, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv4) drop4 = Dropout(0.5)(conv4) pool4 = MaxPooling2D(pool_size=(2, 2))(drop4) conv5 = Conv2D(1024, 3, activation='relu', padding='same', kernel_initializer='he_normal')(pool4) conv5 = Conv2D(1024, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv5) drop5 = Dropout(0.5)(conv5) # 解码器部分 up6 = Conv2D(512, 2, activation='relu', padding='same', kernel_initializer='he_normal')(UpSampling2D(size=(2, 2))(drop5)) merge6 = concatenate([drop4, up6], axis=3) conv6 = Conv2D(512, 3, activation='relu', padding='same', kernel_initializer='he_normal')(merge6) conv6 = Conv2D(512, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv6) up7 = Conv2D(256, 2, activation='relu', padding='same', kernel_initializer='he_normal')(UpSampling2D(size=(2, 2))(conv6)) merge7 = concatenate([conv3, up7], axis=3) conv7 = Conv2D(256, 3, activation='relu', padding='same', kernel_initializer='he_normal')(merge7) conv7 = Conv2D(256, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv7) up8 = Conv2D(128, 2, activation='relu', padding='same', kernel_initializer='he_normal')(UpSampling2D(size=(2, 2))(conv7)) merge8 = concatenate([conv2, up8], axis=3) conv8 = Conv2D(128, 3, activation='relu', padding='same', kernel_initializer='he_normal')(merge8) conv8 = Conv2D(128, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv8) up9 = Conv2D(64, 2, activation='relu', padding='same', kernel_initializer='he_normal')(UpSampling2D(size=(2, 2))(conv8)) merge9 = concatenate([conv1, up9], axis=3) conv9 = Conv2D(64, 3, activation='relu', padding='same', kernel_initializer='he_normal')(merge9) conv9 = Conv2D(64, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv9) conv9 = Conv2D(2, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv9) # 输出层 outputs = Conv2D(1, 1, activation='sigmoid')(conv9) model = Model(inputs=inputs, outputs=outputs) return model ``` 5. 编写训练脚本: ```python import tensorflow as tf from tensorflow.keras.preprocessing.image import ImageDataGenerator from tensorflow.keras.callbacks import ModelCheckpoint, EarlyStopping # 加载数据 train_datagen = ImageDataGenerator( rescale=1./255, shear_range=0.2, zoom_range=0.2, horizontal_flip=True) train_generator = train_datagen.flow_from_directory( 'train/', target_size=(256, 256), batch_size=16, class_mode='binary') # 构建模型 model = Unet((256, 256, 3)) model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy']) # 设置回调函数 checkpoint = ModelCheckpoint('model.h5', monitor='val_loss', save_best_only=True, verbose=1) early_stop = EarlyStopping(monitor='val_loss', patience=10, verbose=1) # 训练模型 model.fit(train_generator, epochs=50, validation_data=val_generator, callbacks=[checkpoint, early_stop]) ``` 6. 模型评估:使用测试集对模型进行评估,并根据评估结果进行调整。 ```python val_datagen = ImageDataGenerator(rescale=1./255) val_generator = val_datagen.flow_from_directory( 'val/', target_size=(256, 256), batch_size=16, class_mode='binary') model.evaluate(val_generator) ``` 7. 模型应用:将训练好的模型用于实际应用中。 ```python from PIL import Image import numpy as np model = tf.keras.models.load_model('model.h5') img = Image.open('test.jpg') img = img.resize((256, 256)) img_array = np.array(img) img_array = np.expand_dims(img_array, axis=0) pred = model.predict(img_array) pred = np.squeeze(pred) pred = np.where(pred > 0.5, 1, 0) output = Image.fromarray(np.uint8(pred * 255)) output.show() ``` 以上是基于自己的数据训练Unet模型的步骤,您可以根据实际情况进行调整和修改。如果您有其他问题,可以随时问我。
评论 9
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包

打赏作者

听风吹等浪起

你的鼓励将是我创作的最大动力

¥1 ¥2 ¥4 ¥6 ¥10 ¥20
扫码支付:¥1
获取中
扫码支付

您的余额不足,请更换扫码支付或充值

打赏作者

实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值