import os
# 图片存储在当前路径(os.getcwd())下,data文件夹中的test文件夹中
input_dir = os.path.join(os.getcwd(), "data", "test")
output_dir = os.path.join(os.getcwd(), "data", "result")
将test中的图片按照名字排序并且存储到一个list当中去——等下图片的读取就会按照这个list进行
import os
"""
example:color_dsift_flower_1.png
item 是一个字符串,item = color_dsift_flower_1
这里千万注意我们只要图片的名字而不要后缀(.png)
images_name是一个存储图片名字的list
"""
images_name = sorted(list({for item[:-4] in os.listdir(input_dir)}))
特别注意:
skimage.imread(fp)返回的是numpy数据,h*w*3,正常RGB通道显示,但在pytorch中我们要将其转换成b*c*w*h的形式
from skimage import io
for image_name in images_name:
img1 = io.imread(os.path.join(input_dir, image_name + ".png"))
同理:opencv获得的numpy的数据顺序为h*w*c,其中张量的最后一个维度才是通道,所以在送入到torch中之前,需要对张量的维度顺序进行变换
重点问题:opencv读取图片是BGR的顺序,我们要将其转换成RGB的形式
import cv2 as cv
for image_name in images_name:
img1 = cv.imread(os.path.join(input_dir, image_name)) # 默认是BRG
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # 将其转换成RGB
你的测试集里面可能既有彩色图片又有灰度图片这时候就要做出判断了——这种情况比较少
ndim = img1.ndim # 这样就知道了img1的维度
if ndim == 2: # 是灰度图了你要怎么办
...
else: # 是彩色图你要怎么办
...
这里进行下变化是为了方便后续图片的预处理,也可以不用的,看你自己——建议用
参考连接
import PIL.Image
# 实现了一个array到image的转换
img1_pil = PIL.Image.fromarray(img1)
拓展:PIL image转换成array
img = np.asarray(image)
pytorch中输入网络中的都是b*c*w*h,这个很重要,参考链接
重点:这个测试transforms.Compose要与你训练的transforms.Compose相同
官方文档
import torchvision.transforms as transforms
"""
transforms.Compose()是容器,里面的是方法
-----------------------------------------------------------------------------------
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
是对PILImage进行的,分别对其进行随机大小和随机宽高比的裁剪,之后resize到指定大小224
------------------------------------------------------------------------------------
transforms.ToTensor(),
把一个取值范围是[0,255]的PIL.Image或者shape为(H,W,C)的numpy.ndarray,转换成形状为[C,H,W],取值范围是[0,1.0]的torch.FloadTensor
-------------------------------------------------------------------------------------------
transforms.Normalize()
用给定的均值和标准差分别对每个通道的数据进行正则化。
"""
# 实例化
data_transforms = transforms.Compose([
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
# 处理图片
# data_transforms处理过后是c*w*h,还少了batch_size这个维度,所以unsqueeze(0)给他加上了这个维度
device = torch.device('cpu') # 有gpu可以改成cuda
img_tensor = (data_transforms(img2_gray_pil)).unsqueeze(0).to(device).float()
device = torch.device('cpu') # cuda 或者 cuda:0
model = Net() # 实例化,Net是你自己的网络的名字,在后续train的部分会告诉你
# 参数文件的所在路径
model_path = os.path.join(os.getcwd(), "nets", "parameters","xxxxxxxx.pkl")
# 将参数传递到网络中去
model.load_state_dict(torch.load(self.model_path, map_location=self.device))
# 将模型放到cpu或者gpu中
model.to(device)
"""
在PyTorch中进行test时,会使用model.eval()切换到测试模式
with torch.no_grad()则主要是用于停止autograd模块的工作,以起到加速和节省显存的作用。它的作用是将该with语句包裹起来的部分停止梯度的更新,从而节省了GPU算力和显存,但是并不会影响dropout和BN层的行为。
如果不在意显存大小和计算时间的话,仅仅使用model.eval()已足够得到正确的validation/test的结果;而with torch.no_grad()则是更进一步加速和节省gpu空间(因为不用计算和存储梯度),从而可以更快计算,也可以跑更大的batch来测试。
总结:用了model.eval()后就可以不用with torch.no_grad()
"""
model.eval()
拓展链接:关于model.train()、model.eval()、torch.no_grad()
class Net(nn.Module):
def __init__(self):
super(SESFuseNet, self).__init__()
'''
在这里进行一些块的实例化,比如注意力模块
'''
@staticmethod
def conv_block(in_channels, out_channels, kernel_size=3):
"""
The conv block of common setting: conv -> relu -> bn
如果需要池化的话可以自行添加,我上次没用到所以不加了
In conv operation, the padding = 1
:param in_channels: int, the input channels of feature
:param out_channels: int, the output channels of feature
:param kernel_size: int, the kernel size of feature
:return:
"""
block = torch.nn.Sequential(
torch.nn.Conv2d(kernel_size=kernel_size, in_channels=in_channels, out_channels=out_channels, padding=1),
torch.nn.ReLU(),
torch.nn.BatchNorm2d(out_channels),
)
return block
@staticmethod
def concat(f1, f2):
"""
Concat two feature in channel direction
"""
return torch.cat((f1, f2), 1)
def forward(self, img1):
'''
具体的网络实现
'''
return ...
由于不同的任务有不同的损失函数,所以在这里便不举例了
experiment_name = 'xxx' # 参数文件的命名
gpu_device = "cuda:0"
learning_rate = 1e-4
epochs = 30
batch_size = 1
shuffle = True # 是否随机读取数据集与测试集
3.2、各种地址
# address
project_addrsss = os.getcwd()
train_dir = os.path.join(project_addrsss, "data", "coco2014", "train2014") # 训练集地址
val_dir = os.path.join(project_addrsss, "data", "coco2014", "val2014") # 测试集地址
log_address = os.path.join(project_addrsss, "nets", "train_record", experiment_name + "_log_file.txt") # 日志的保存地址
is_out_log_file = True # 等下查
parameter_address = os.path.join(project_addrsss, "nets", "parameters") # 最终的参数保存地址
3.3、
# 图片预处理
data_transforms = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize([0.4500517361627943], [0.26465333914691797]),
])
3.4、训练中的难点——dataloader
单独写了一个博客介绍
看完上面博客在看下面
'''用一个字典来存储训练以及测试tensor'''
image_datasets = {}
'''COCODataset这个类在上面那个博客里面'''
image_datasets['train'] = COCODataset(train_dir, transform=data_transforms, need_crop=False, need_augment=False)
image_datasets['val'] = COCODataset(val_dir, transform=data_transforms, need_crop=False, need_augment=False)
dataloders = {}
'''将上述tensor放到DataLoader中去'''
dataloders['train'] = DataLoader(
image_datasets['train'],
batch_size=batch_size,
shuffle=shuffle,
num_workers=1)
dataloders['val'] = DataLoader(
image_datasets['val'],
batch_size=batch_size,
shuffle=shuffle,
num_workers=1)
'''日志的输出以及存储'''
datasets_sizes = {x: len(image_datasets[x]) for x in ['train', 'val']} # 数据集中的图片数量
print_and_log("datasets size: {}".format(datasets_sizes), is_out_log_file, log_address)
这里只介绍了图片的读取,lable的读取也是一样的
3.5、定义训练函数
3.5.1、基础版本,无任何画图、保存日志等功能
model = Net() # 网络实例化
model.to(device)
criterion = Loss().to(device) # 损失函数实例化,并且放到device上运算(loss是你自己定义的函数)
optimizer = optim.Adam(model.parameters(), learning_rate) # 实例化优化器
def train(epoch): # 这里的epoch代表的是第几次迭代
model.train() # 作用是启用batch normalization和drop out。
running_loss = 0.0 # 将每个batch_size的损失加到一起
for X, y in enumerate(dataloders['train']):
input = X.to(device) # 将数据拷贝到GPU上
output = model.forward(X) # 将数据前向传入训练模型
loss = criterion(output , y)
running_loss += loss.item()
# 这一部分永远是固定的
optimizer.zero_grad()
loss.backward()
optimizer.step()
epoch_loss_train = running_loss / datasets_sizes # datasets_sizes是你数据集的大小
return epoch_loss_train
3.6、测试——每个epoch过后就做一次测试
def val():
model.eval() # 模型验证——关闭drop out
running_loss = 0.0
with torch.no_grad(): # 不用梯度了
for i, data in enumerate(dataloders['val']):
input = data.to(gpu_device) # 将数据拷贝到GPU上
optimizer.zero_grad() # 梯度置零
output = model.forward('train', input)
loss, lp_loss, lssim_loss = criterion(image_in=input, image_out=output)
running_loss += loss.item()
epoch_loss_val = running_loss / datasets_sizes # datasets_sizes验证集数据的大小
return epoch_loss_val
3.7mian()
def main():
min_loss = 100000000.0 # 保存最好的模型
loss_train = [] # 存储每个epoch的损失
loss_val = [] # 存储每个epoch之后的训练结果
since = time.time() # 记录时间
for epoch in range(epochs):
epoch_loss_train = train(epoch)
loss_train.append(epoch_loss_train)
epoch_loss_val = val()
loss_val.append(epoch_loss_val)
'''如果测试的结果很好就将其保存下来'''
if epoch_loss_val < min_loss:
min_loss = epoch_loss_val
best_model_wts = model.state_dict()
torch.save(best_model_wts,os.path.join(parameter_address, experiment_name + '.pkl')) # 保存最好的数据模型
# 每个epoch的训练结果都保存下来
model_wts = model.state_dict() # 保存训练模型
torch.save(model_wts,os.path.join(parameter_address, experiment_name + "_" + str(epoch) + '.pkl')) # 保存模型参数
time_elapsed = time.time() - since # 获取程序运行时间
日志输出以及图片的常用工具
import os
import cv2
import numpy as np
from IPython.display import clear_output
import matplotlib.pyplot as plt
import torch
import random
# 设置随机数种子,为了成功复现实验
# 传入的数值用于指定随机数生成时所用算法开始时所选定的整数值,如果使用相同的seed()值,则每次生成的随机数都相同;
# 如果不设置这个值,则系统会根据时间来自己选择这个值,此时每次生成的随机数会因时间的差异而有所不同。
def training_setup_seed(seed):
# 设置随机种子后,是每次运行test.py文件的输出结果都一样,而不是每次随机函数生成的结果一样:
torch.manual_seed(seed) # 设置CPU生成随机数的种子,返回一个torch.Generator对象
torch.cuda.manual_seed_all(seed) # 为所有的GPU设置种子
np.random.seed(seed)
random.seed(seed)
torch.backends.cudnn.deterministic = True # 设置为True时,每次返回的卷积算法是确定的,即默认算法。
# 如果配合上设置 torch 的随机种子为固定值,可以保证每次运行网络的时候相同输入的输出是固定的。
# 调整学习率
def adjust_learning_rate(optimizer, learning_rate, epoch):
"""Sets the learning rate to the initial LR decayed by half every 10 epochs until 1e-5"""
lr = learning_rate * (0.8 ** (epoch // 2)) # 每两个epoch,学习率乘上0.8。 lr=learning_rate(0.8^(epoch // 2))
# if not lr < 1e-6:
for param_group in optimizer.param_groups: # 动态修改学习率
param_group['lr'] = lr
# 绘制最终结果损失函数图
def plot_loss(experiment_name, epoch, train_loss_list, val_loss_list):
clear_output(True)
# train_loss_list[-1]返回train_loss_list中倒数第一个值
print('Epoch %s. train loss: %s. val loss: %s' % (epoch, train_loss_list[-1], val_loss_list[-1]))
print('Best val loss: %s' % (min(val_loss_list)))
print('Back up')
print('train_loss_list: {}'.format(train_loss_list))
print('val_loss_list: {}'.format(val_loss_list))
plt.figure()
plt.plot(train_loss_list, color="r", label="train loss")
plt.plot(val_loss_list, color="b", label="val loss")
plt.legend(loc="best") # 设置图例位置
plt.xlabel("Epochs") # 一个epoch包含数个Iteration
plt.ylabel("Loss")
plt.title("Loss " + experiment_name, fontsize=16)
figure_address = os.path.join(os.path.join(os.getcwd(), 'nets'), 'figure') # 最后结果图的保存位置 ./nets/figure
plt.savefig(os.path.join(figure_address, experiment_name + '_loss')) # 保存的图片名称
plt.show()
# 绘制迭代过程的损失函数图
def plot_iteration_loss(experiment_name, epoch, loss, lp_loss, lssim_loss): #epocah
# clear_output(True)
# print('Iteration %s. loss: %s.' % (iteration, loss))
# print('Iteration %s. lp_loss: %s.' % (iteration, lp_loss))
# print('Iteration %s. lssim_loss: %s.' % (iteration, lssim_loss))
# print('loss: {}'.format(loss))
# print('lp_loss: {}'.format(lp_loss))
# print('lssim_loss: {}'.format(lssim_loss))
plt.figure()
plt.plot(loss, color="r", label="loss")
plt.plot(lp_loss, color="g", label="lp_loss")
plt.plot(lssim_loss, color="b", label="lssim_loss")
plt.legend(loc="best")
plt.xlabel("Iterations") # Iterations(迭代次数) 等于 数据集拆分的batch数。一个batch的数据用来为模型计算一次梯度下降更新
plt.ylabel("Loss")
plt.title("Loss " + experiment_name, fontsize=16)
figure_address = os.path.join(os.path.join(os.getcwd(), 'nets'), 'figures') # 迭代过程的损失函数结果图
plt.savefig(os.path.join(figure_address, experiment_name + '_' + str(epoch) + '_loss')) # epocah
plt.show(block=False)
# 生成训练记录 log_file文件
def print_and_log(content, is_out_log_file=True, file_address=None):
print(content)
if is_out_log_file:
f = open(os.path.join(file_address), "a") # join()是将多个路径和文件组成新的路径
f.write(content)
f.write("\n")
f.close()
# 计算像素平均值
def get_mean_value(input_dir): # input_dir是./data/multi_focus文件夹
# os.path.join()拼接待操作对象input_dir,item os.listdir()返回输入路径下的文件和列表名称作为list的元素
images_list = [os.path.join(input_dir, item) for item in sorted(os.listdir(input_dir))] # sort()用来排序
count = 0
pixel_sum = 0 # 初始化像素和
for index, sub_folder in enumerate(images_list):
image_name = os.path.basename(sub_folder)
last_image = cv2.imread(os.path.join(sub_folder, image_name + "_1.png"), 0) * 1.0 / 255
next_image = cv2.imread(os.path.join(sub_folder, image_name + "_2.png"), 0) * 1.0 / 255
pixel_sum = pixel_sum + np.sum(last_image) + np.sum(next_image) # 计算像素和
count = count + last_image.size + next_image.size # 计数
return pixel_sum / count # 求平均
def get_std_value(input_dir, mean):
images_list = [os.path.join(input_dir, item) for item in sorted(os.listdir(input_dir))]
count = 0 # 初始化
pixel_sum = 0 # 初始化像素和
for index, sub_folder in enumerate(images_list): # enumerate()用来遍历集合对象,同时还能得到当前元素的索引值
image_name = os.path.basename(sub_folder)
# np.power(a,b) 计算a的b次方
# 计算像素与平均值之间的方差
last_image = np.power((cv2.imread(os.path.join(sub_folder, image_name + "_1.png"), 0) * 1.0 / 255) - mean, 2)
next_image = np.power((cv2.imread(os.path.join(sub_folder, image_name + "_2.png"), 0) * 1.0 / 255) - mean, 2)
pixel_sum = pixel_sum + np.sum(last_image) + np.sum(next_image)
count = count + last_image.size + next_image.size
return np.sqrt(pixel_sum / count) # 开根号