|
- '''
- # @time:2023/2/26 9:49
- # Author:Tuan
- # @File:train_memory.py
- '''
- '''
- # @time:2022/11/27 21:19
- # Author:Tuan
- # @File:train_sample.py
- '''
- import os.path
- import numpy as np
- import torch
- import torch.nn
- from torch.nn import functional as F
- from torch.nn import CrossEntropyLoss
- from torch.utils.data import DataLoader
- from utils.dataset.TT_Dataset import MyDataset # 读取数据所用函数
- # from Model.HrNet.hrnet_M import HighResolutionNet
- from Model.UNet.u_r_m_3l import UNet
- import imageio
-
- device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
- def estimate(y_label, y_pred):
- # y_pred[y_label==0]=0
- # 准确率
- acc = np.mean(np.equal(y_label.cpu().numpy(), y_pred.cpu().numpy()) + 0)
-
- return acc, y_pred
-
- def model_predict(model, img_data, lab_data, img_size):
- row, col, dep = img_data.shape
-
- if row % img_size != 0 or col % img_size != 0:
- # 计算填充后图像的 hight 和 width
- padding_h = (row // img_size + 1) * img_size
- padding_w = (col // img_size + 1) * img_size
- else:
- # 不填充后图像的 hight 和 width
- padding_h = (row // img_size) * img_size
- padding_w = (col // img_size) * img_size
-
- # 初始化一个 0 矩阵,将图像的值赋值到 0 矩阵的对应位置
- padding_img = np.zeros((padding_h, padding_w, dep), dtype='float32')
- padding_img[:row, :col, :] = img_data[:row, :col, :]
-
- # 初始化一个 0 矩阵,用于将预测结果的值赋值到 0 矩阵的对应位置
- padding_pre = np.zeros((padding_h, padding_w), dtype='uint8')
- padding_pre = torch.tensor(padding_pre)
- padding_pre = padding_pre.to(device)
-
- # 对 img_size * img_size 大小的图像进行预测
- count = 0 # 用于计数
- for i in list(np.arange(0, padding_h, img_size)):
- if (i + img_size) > padding_h:
- continue
- for j in list(np.arange(0, padding_w, img_size)):
- if (j + img_size) > padding_w:
- continue
-
- # 取 img_size 大小的图像,在第一维添加维度,变成四维张量,用于模型预测
- img_data_ = padding_img[i:i + img_size, j:j + img_size, :]
- img_data_ = img_data_[np.newaxis, :, :, :]
- img_data_ = np.transpose(img_data_, (0, 3, 1, 2))
- img_data_ = torch.from_numpy(img_data_)
-
- # 预测,对结果进行处理
- y_pre, _ = model(img_data_.to(device),1)
- # y_pre = model.predict(img_data_)
- y_pre = np.squeeze(y_pre, axis=0)
- y_pre = torch.argmax(y_pre, axis=0)
- # y_pre = y_pre.astype('uint8')
-
- # 将预测结果的值赋值到 0 矩阵的对应位置
- padding_pre[i:i + img_size, j:j + img_size] = y_pre[:img_size, :img_size]
-
- count += 1 # 每预测一块就+1
-
- # 计算准确率
- acc = estimate(lab_data.to(device), padding_pre[:row, :col]+1)
- return acc
-
- #超参数
- num_classes = 7
- batch_size = 1
- epoches = 200
- # modelname = "UNet"
- #读取数据
- imagePath = r"E:\yqj\try\code\torch\Train\Data\coastline\img"
- labelPath = r"E:\yqj\try\code\torch\Train\Data\coastline\lab_type"
- # lab_sl = r"D:\code\torch\Train\coastline\lab_SL"
- # lab_line = r"D:\code\torch\Train\coastline\lab_line"
- trainDataset = MyDataset(imagePath, labelPath)
- trainDatasetloader = DataLoader(trainDataset, batch_size, shuffle=True)
- trainLen = len(trainDatasetloader)
-
- # test_image = imageio.imread(r"D:\code\torch\Train\coastline\val\pre_image_vaild\GF1_PMS1_E119.2_N34.7_20160818_L3A0001770062-MSS1.tif")
- # test_label = imageio.imread(r"D:\code\torch\Train\coastline\val\coastline_type_vaild\GF1_PMS1_E119.2_N34.7_20160818_L3A0001770062-MSS1.tif")
- # test_image = torch.from_numpy(test_image)
- # test_label = torch.from_numpy(test_label)
- #构建模型、优化器、损失
- model = UNet(num_classes=7)
- total = sum([param.nelement() for param in model.parameters()])
- print("Number of parameter: %.2fM" % (total / 1e6))
- Loss_type = CrossEntropyLoss()
- Loss_sl = CrossEntropyLoss()
- Loss_line = CrossEntropyLoss()
- optimizer = torch.optim.AdamW(model.parameters(), lr=0.001, weight_decay=0.01, betas=(0.9, 0.95))
- #开始训练
- for epoch in range(epoches):
- print(
- f"\n----------------------------------------------epoch: {epoch}----------------------------------------------")
- loss_total = 0
- for i,data in enumerate(trainDatasetloader):
- # if i == 2:
- # break
- image, label = data
- image = image.to(device)
- label = label.to(device)
- # lab_sl = lab_sl.to(device)
- # lab_line = lab_line.to(device)
-
- #梯度清零
- optimizer.zero_grad()
- m = ((1 - epoch / epoches) ** 0.9) * (0.9 - 0.9 / 100) + 0.9 / 100
- # output , output_sl, output_line, memory_input = model(image)
- output, memory_input = model(image)
- with torch.no_grad():
- model.update(
- # features = memory_input,
- features=F.interpolate(memory_input, size=(256,256), mode='bilinear'),
- lab=label,m=m
- )
- loss_type = Loss_type(output, label.long())
- # loss_sl = Loss_sl(output_sl, lab_sl.long())
- # loss_line = Loss_line(output_line, lab_line.long())
- laph = 1.2
- loss_step = loss_type #+ loss_sl + laph * loss_line
- print("\r train: epoch: {}, step: {}/{}, loss: {} ".format(epoch, i, trainLen,loss_step ,end=''))
- loss_step.backward()
- optimizer.step()
- loss_total += loss_step
-
- loss_epoch = loss_total/trainLen
- print("\r epoch: {}, epoch_loss: {}".format(epoch, round(float(loss_epoch), 8)), end='')
- # acc, pre = model_predict(model, test_image, test_label, 256)
- # print(f'acc:{acc}')
-
- #保存模型
- # Path = r"D:\code\torch\Model\memory\model\test.pth"
- # torch.save(model.state_dict(), Path)
|