from torch.utils import data
import os
from os.path import join, abspath, splitext, split, isdir, isfile
from PIL import Image
import numpy as np
import cv2
这段代码导入了一些常用的Python库和模块，用于数据处理和图像处理。让我逐个解释这些导入语句的作用：
from torch.utils import data
这行代码导入了torch.utils.data模块，该模块提供了用于数据加载和处理的工具，特别是在使用PyTorch进行深度学习时。它包含了一些用于创建自定义数据集和数据加载器的类和函数。
import os
这行代码导入了Python的os模块，用于与操作系统进行交互，例如访问文件和目录。
from os.path import join, abspath, splitext, split, isdir, isfile
这行代码从os.path模块中导入了一些函数，包括join、abspath、splitext、split、isdir和isfile。这些函数用于处理文件路径和判断路径是否为目录或文件。
from PIL import Image
这行代码导入了Python Imaging Library（PIL）的Image模块，用于图像的读取、处理和保存。
import numpy as np
这行代码导入了Python的numpy库，用于进行数值计算和数组操作。numpy提供了高效的多维数组对象和相关函数，常用于科学计算和数据处理。
import cv2
这行代码导入了OpenCV库（Open Source Computer Vision Library）的Python接口。OpenCV是一个广泛使用的计算机视觉库，提供了各种图像和视频处理功能，包括读取、显示、保存图像，以及图像处理和计算机视觉算法的实现。
通过导入这些库和模块，你可以在代码中使用它们提供的函数和类来处理数据和图像，进行深度学习和计算机视觉任务。
def prepare_image_PIL(im):
im = im[:,:,::-1] - np.zeros_like(im) # rgb to bgr
im -= np.array((104.00698793,116.66876762,122.67891434))
im = np.transpose(im, (2, 0, 1)) # (H x W x C) to (C x H x W)
return im
这段代码定义了一个名为prepare_image_PIL的函数，该函数接受一个图像作为输入，并对其进行预处理。让我逐行解释代码的作用：
im = im[:,:,::-1] - np.zeros_like(im)
这行代码将图像的颜色通道从RGB顺序转换为BGR顺序。它通过使用[:,:,::-1]对图像的颜色通道进行反转，将RGB通道顺序变为BGR通道顺序。然后，通过减去np.zeros_like(im)，将图像减去一个与图像形状相同的全零数组，实现了图像通道顺序的转换。
im -= np.array((104.00698793,116.66876762,122.67891434))
这行代码对图像进行减去均值的操作。它从图像中减去了一个均值数组(104.00698793, 116.66876762, 122.67891434)，这些值是在训练数据集上计算得到的图像均值。这种减去均值的操作常用于图像预处理，以减少输入数据的冗余信息。
im = np.transpose(im, (2, 0, 1))
这行代码对图像的维度进行转置操作，将通道维度从最后一维移动到第一维。它使用np.transpose函数和(2, 0, 1)作为参数，表示将维度顺序从(H, W, C)转换为(C, H, W)，其中H表示图像的高度，W表示图像的宽度，C表示图像的通道数。
最后，函数返回经过预处理的图像im。
这个函数的目的是将输入的图像进行预处理，以便用于某些深度学习模型或计算机视觉任务。预处理操作包括颜色通道顺序转换、减去均值和维度转置。
def prepare_image_cv2(im):
im -= np.array((104.00698793,116.66876762,122.67891434))
im = np.transpose(im, (2, 0, 1)) # (H x W x C) to (C x H x W)
return im
这段代码定义了一个名为prepare_image_cv2的函数，该函数接受一个图像作为输入，并对其进行预处理。让我逐行解释代码的作用：
im -= np.array((104.00698793,116.66876762,122.67891434))
这行代码对图像进行减去均值的操作。它从图像中减去了一个均值数组(104.00698793, 116.66876762, 122.67891434)，这些值是在训练数据集上计算得到的图像均值。这种减去均值的操作常用于图像预处理，以减少输入数据的冗余信息。
im = np.transpose(im, (2, 0, 1))
这行代码对图像的维度进行转置操作，将通道维度从最后一维移动到第一维。它使用np.transpose函数和(2, 0, 1)作为参数，表示将维度顺序从(H, W, C)转换为(C, H, W)，其中H表示图像的高度，W表示图像的宽度，C表示图像的通道数。
最后，函数返回经过预处理的图像im。
这个函数的目的是将输入的图像进行预处理，以便用于某些深度学习模型或计算机视觉任务。预处理操作包括减去均值和维度转置。与prepare_image_PIL函数相比，这个函数没有进行颜色通道顺序转换。这是因为cv2模块读取的图像默认为BGR顺序，而不是RGB顺序。
class BSDS_RCFLoader(data.Dataset):
"""
Dataloader BSDS500
"""
def __init__(self, root='data/HED-BSDS_PASCAL', split='train', transform=False):
self.root = root
self.split = split
self.transform = transform
if self.split == 'train':
self.filelist = join(self.root, 'train_pair.lst')
elif self.split == 'test':
self.filelist = join(self.root, 'test.lst')
else:
raise ValueError("Invalid split type!")
with open(self.filelist, 'r') as f:
self.filelist = f.readlines()
def __len__(self):
return len(self.filelist)
def __getitem__(self, index):
if self.split == "train":
img_file, lb_file = self.filelist[index].split()
lb = np.array(Image.open(join(self.root, lb_file)), dtype=np.float32)
if lb.ndim == 3:
lb = np.squeeze(lb[:, :, 0])
assert lb.ndim == 2
lb = lb[np.newaxis, :, :]
lb[lb == 0] = 0
lb[np.logical_and(lb>0, lb<127.5)] = 2
lb[lb >= 127.5] = 1
else:
img_file = self.filelist[index].rstrip()
if self.split == "train":
img = np.array(cv2.imread(join(self.root, img_file)), dtype=np.float32)
img = prepare_image_cv2(img)
return img, lb
else:
img = np.array(Image.open(join(self.root, img_file)), dtype=np.float32)
img = prepare_image_PIL(img)
return img
这段代码定义了一个名为BSDS_RCFLoader的类，该类是一个自定义的数据加载器，用于加载BSDS500数据集。让我逐行解释代码的作用：
def __init__(self, root='data/HED-BSDS_PASCAL', split='train', transform=False):
self.root = root
self.split = split
self.transform = transform
这是类的构造函数，用于初始化对象的属性。它接受三个参数：root表示数据集的根目录，默认为data/HED-BSDS_PASCAL；split表示数据集的划分类型，默认为train；transform表示是否进行数据转换，默认为False。
if self.split == 'train':
self.filelist = join(self.root, 'train_pair.lst')
elif self.split == 'test':
self.filelist = join(self.root, 'test.lst')
else:
raise ValueError("Invalid split type!")
根据split的值，确定要加载的数据集划分类型。如果split为train，则设置filelist为root目录下的train_pair.lst文件的路径；如果split为test，则设置filelist为root目录下的test.lst文件的路径。如果split既不是train也不是test，则抛出一个ValueError异常。
with open(self.filelist, 'r') as f:
self.filelist = f.readlines()
使用open函数打开filelist文件，并以只读模式读取其中的内容。然后，使用readlines方法将文件内容按行读取，并将结果赋值给self.filelist。这样，self.filelist就成为一个包含文件中每一行内容的列表。
这个类的目的是作为BSDS500数据集的数据加载器，根据给定的划分类型加载相应的文件列表。这个类的实例可以用于迭代访问数据集中的样本，并根据需要进行数据转换。
这段代码定义了BSDS_RCFLoader类的两个特殊方法：__len__和__getitem__。这些方法用于支持类的实例像列表一样进行索引和迭代。
def __len__(self):
return len(self.filelist)
__len__方法返回数据集中样本的数量，即self.filelist列表的长度。
def __getitem__(self, index):
if self.split == "train":
img_file, lb_file = self.filelist[index].split()
lb = np.array(Image.open(join(self.root, lb_file)), dtype=np.float32)
if lb.ndim == 3:
lb = np.squeeze(lb[:, :, 0])
assert lb.ndim == 2
lb = lb[np.newaxis, :, :]
lb[lb == 0] = 0
lb[np.logical_and(lb>0, lb<127.5)] = 2
lb[lb >= 127.5] = 1
else:
img_file = self.filelist[index].rstrip()
__getitem__方法根据给定的索引index返回对应的样本。如果数据集划分类型为train，则从self.filelist中获取索引index对应的图像文件名img_file和标签文件名lb_file。然后，使用Image.open函数打开标签图像文件，并将其转换为np.array类型的数组lb，数据类型为np.float32。如果标签图像的维度为3，则使用np.squeeze函数将其压缩为2维。然后，确保标签图像的维度为2。接下来，对标签图像进行一些操作，将0值保持为0，将介于0和127.5之间的值设为2，将大于等于127.5的值设为1。
如果数据集划分类型不是train，则直接从self.filelist中获取索引index对应的图像文件名img_file。
if self.split == "train":
img = np.array(cv2.imread(join(self.root, img_file)), dtype=np.float32)
img = prepare_image_cv2(img)
return img, lb
else:
img = np.array(Image.open(join(self.root, img_file)), dtype=np.float32)
img = prepare_image_PIL(img)
return img
根据数据集划分类型，加载对应的图像文件并进行预处理。如果划分类型为train，则使用cv2.imread函数读取图像文件，并将其转换为np.array类型的数组img，数据类型为np.float32。然后，使用prepare_image_cv2函数对图像进行预处理。最后，返回预处理后的图像img和标签图像lb。
如果划分类型不是train，则使用Image.open函数打开图像文件，并将其转换为np.array类型的数组img，数据类型为np.float32。然后，使用prepare_image_PIL函数对图像进行预处理。最后，只返回预处理后的图像img。
这些方法的目的是根据索引获取数据集中的样本，并根据需要进行图像的加载和预处理。在训练阶段，返回图像和对应的标签；在测试阶段，只返回图像。
import numpy as np
import torch
from torch.optim.optimizer import Optimizer, required
def cross_entropy_loss_RCF(prediction, label):
label = label.long()
mask = label.float()
num_positive = torch.sum((mask==1).float()).float()
num_negative = torch.sum((mask==0).float()).float()
mask[mask == 1] = 1.0 * num_negative / (num_positive + num_negative)
mask[mask == 0] = 1.1 * num_positive / (num_positive + num_negative)
mask[mask == 2] = 0
cost = torch.nn.functional.binary_cross_entropy(
prediction.float(),label.float(), weight=mask, reduce=False)
return torch.sum(cost)
class SGD_caffe(Optimizer):
"""Implements stochastic gradient descent (optionally with momentum).
Nesterov momentum is based on the formula from
`On the importance of initialization and momentum in deep learning`__.
Args:
params (iterable): iterable of parameters to optimize or dicts defining
parameter groups
lr (float): learning rate
momentum (float, optional): momentum factor (default: 0)
weight_decay (float, optional): weight decay (L2 penalty) (default: 0)
dampening (float, optional): dampening for momentum (default: 0)
nesterov (bool, optional): enables Nesterov momentum (default: False)
Example:
>>> optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9)
>>> optimizer.zero_grad()
>>> loss_fn(model(input), target).backward()
>>> optimizer.step()
__ http://www.cs.toronto.edu/%7Ehinton/absps/momentum.pdf
.. note::
The implementation of SGD with Momentum/Nesterov subtly differs from
Sutskever et. al. and implementations in some other frameworks.
Considering the specific case of Momentum, the update can be written as
.. math::
v = \rho * v + g \\
p = p - lr * v
where p, g, v and :math:`\rho` denote the parameters, gradient,
velocity, and momentum respectively.
This is in contrast to Sutskever et. al. and
other frameworks which employ an update of the form
.. math::
v = \rho * v + lr * g \\
p = p - v
The Nesterov version is analogously modified.
"""
def __init__(self, params, lr=required, momentum=0, dampening=0,
weight_decay=0, nesterov=False):
if lr is not required and lr < 0.0:
raise ValueError("Invalid learning rate: {}".format(lr))
if momentum < 0.0:
raise ValueError("Invalid momentum value: {}".format(momentum))
if weight_decay < 0.0:
raise ValueError("Invalid weight_decay value: {}".format(weight_decay))
defaults = dict(lr=lr, momentum=momentum, dampening=dampening,
weight_decay=weight_decay, nesterov=nesterov)
if nesterov and (momentum <= 0 or dampening != 0):
raise ValueError("Nesterov momentum requires a momentum and zero dampening")
super(SGD_caffe, self).__init__(params, defaults)
def __setstate__(self, state):
super(SGD_caffe, self).__setstate__(state)
for group in self.param_groups:
group.setdefault('nesterov', False)
def step(self, closure=None):
"""Performs a single optimization step.
Arguments:
closure (callable, optional): A closure that reevaluates the model
and returns the loss.
"""
loss = None
if closure is not None:
loss = closure()
for group in self.param_groups:
weight_decay = group['weight_decay']
momentum = group['momentum']
dampening = group['dampening']
nesterov = group['nesterov']
for p in group['params']:
if p.grad is None:
continue
d_p = p.grad.data
if weight_decay != 0:
d_p.add_(weight_decay, p.data)
if momentum != 0:
param_state = self.state[p]
if 'momentum_buffer' not in param_state:
buf = param_state['momentum_buffer'] = torch.zeros_like(p.data)
buf.mul_(momentum).add_(group['lr'], d_p)
else:
buf = param_state['momentum_buffer']
buf.mul_(momentum).add_(group['lr'], d_p)
if nesterov:
d_p = d_p.add(momentum, buf)
else:
d_p = buf
p.data.add_(-1, d_p)
return loss
import numpy as np
import torch
from torch.optim.optimizer import Optimizer, required
def cross_entropy_loss_RCF(prediction, label):
label = label.long()
mask = label.float()
num_positive = torch.sum((mask==1).float()).float()
num_negative = torch.sum((mask==0).float()).float()
mask[mask == 1] = 1.0 * num_negative / (num_positive + num_negative)
mask[mask == 0] = 1.1 * num_positive / (num_positive + num_negative)
mask[mask == 2] = 0
cost = torch.nn.functional.binary_cross_entropy(
prediction.float(),label.float(), weight=mask, reduce=False)
return torch.sum(cost)
这段代码定义了一个名为cross_entropy_loss_RCF的函数，用于计算预测值和标签之间的交叉熵损失。
函数的输入参数是prediction（预测值）和label（标签）。首先，将label转换为long类型，并创建一个与label相同形状的掩码mask，数据类型为float。
接下来，计算标签中正样本和负样本的数量。通过将掩码mask与1进行比较，得到一个布尔类型的掩码，其中1表示对应位置上的元素等于1，0表示不等于1。将这个布尔掩码转换为float类型，并使用torch.sum函数计算其中值为1的元素的总和，得到正样本的数量num_positive。同样地，计算值为0的元素的总和，得到负样本的数量num_negative。
然后，根据正负样本的数量，对掩码mask进行一些操作。将掩码中值为1的元素设为1.0 * num_negative / (num_positive + num_negative)，将掩码中值为0的元素设为1.1 * num_positive / (num_positive + num_negative)，将掩码中值为2的元素设为0。
最后，使用torch.nn.functional.binary_cross_entropy函数计算二进制交叉熵损失。该函数的输入参数包括预测值prediction、标签label和权重weight（即掩码mask）。设置reduce参数为False，以便返回每个样本的损失值。最后，使用torch.sum函数计算所有样本的总损失，并返回该值。
这个函数的作用是计算预测值和标签之间的交叉熵损失，其中通过调整正负样本的权重，可以在不平衡数据集中更好地处理样本的类别不平衡问题。
class SGD_caffe(Optimizer):
"""Implements stochastic gradient descent (optionally with momentum).
Nesterov momentum is based on the formula from
`On the importance of initialization and momentum in deep learning`__.
Args:
params (iterable): iterable of parameters to optimize or dicts defining
parameter groups
lr (float): learning rate
momentum (float, optional): momentum factor (default: 0)
weight_decay (float, optional): weight decay (L2 penalty) (default: 0)
dampening (float, optional): dampening for momentum (default: 0)
nesterov (bool, optional): enables Nesterov momentum (default: False)
Example:
>>> optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9)
>>> optimizer.zero_grad()
>>> loss_fn(model(input), target).backward()
>>> optimizer.step()
__ http://www.cs.toronto.edu/%7Ehinton/absps/momentum.pdf
.. note::
The implementation of SGD with Momentum/Nesterov subtly differs from
Sutskever et. al. and implementations in some other frameworks.
Considering the specific case of Momentum, the update can be written as
.. math::
v = \rho * v + g \\
p = p - lr * v
where p, g, v and :math:`\rho` denote the parameters, gradient,
velocity, and momentum respectively.
This is in contrast to Sutskever et. al. and
other frameworks which employ an update of the form
.. math::
v = \rho * v + lr * g \\
p = p - v
The Nesterov version is analogously modified.
"""
def __init__(self, params, lr=required, momentum=0, dampening=0,
weight_decay=0, nesterov=False):
if lr is not required and lr < 0.0:
raise ValueError("Invalid learning rate: {}".format(lr))
if momentum < 0.0:
raise ValueError("Invalid momentum value: {}".format(momentum))
if weight_decay < 0.0:
raise ValueError("Invalid weight_decay value: {}".format(weight_decay))
defaults = dict(lr=lr, momentum=momentum, dampening=dampening,
weight_decay=weight_decay, nesterov=nesterov)
if nesterov and (momentum <= 0 or dampening != 0):
raise ValueError("Nesterov momentum requires a momentum and zero dampening")
super(SGD_caffe, self).__init__(params, defaults)
def __setstate__(self, state):
super(SGD_caffe, self).__setstate__(state)
for group in self.param_groups:
group.setdefault('nesterov', False)
def step(self, closure=None):
"""Performs a single optimization step.
Arguments:
closure (callable, optional): A closure that reevaluates the model
and returns the loss.
"""
loss = None
if closure is not None:
loss = closure()
for group in self.param_groups:
weight_decay = group['weight_decay']
momentum = group['momentum']
dampening = group['dampening']
nesterov = group['nesterov']
for p in group['params']:
if p.grad is None:
continue
d_p = p.grad.data
if weight_decay != 0:
d_p.add_(weight_decay, p.data)
if momentum != 0:
param_state = self.state[p]
if 'momentum_buffer' not in param_state:
buf = param_state['momentum_buffer'] = torch.zeros_like(p.data)
buf.mul_(momentum).add_(group['lr'], d_p)
else:
buf = param_state['momentum_buffer']
buf.mul_(momentum).add_(group['lr'], d_p)
if nesterov:
d_p = d_p.add(momentum, buf)
else:
d_p = buf
p.data.add_(-1, d_p)
return loss
这段代码定义了一个名为SGD_caffe的类，它是torch.optim.optimizer.Optimizer类的子类，实现了随机梯度下降（SGD）算法（可选带有动量）。
在类的注释中提到，Nesterov动量基于论文On the importance of initialization and momentum in deep learning中的公式。该类的构造函数接受一些参数，包括要优化的参数集合params、学习率lr、动量因子momentum、权重衰减因子weight_decay、动量的阻尼因子dampening和是否启用Nesterov动量的标志nesterov。
类中定义了__init__方法用于初始化类的实例。在这个方法中，首先检查学习率、动量和权重衰减是否为非负值，如果不是，则抛出ValueError。然后，创建一个defaults字典，存储参数的默认值。如果启用了Nesterov动量，并且动量小于等于0或阻尼不等于0，则抛出ValueError。最后，调用父类的__init__方法来初始化优化器。
类中还定义了__setstate__方法，用于设置优化器的状态。在这个方法中，将nesterov设置为False，如果参数组中没有设置nesterov。
类中的step方法执行单次优化步骤。它接受一个可调用对象closure，用于重新评估模型并返回损失值。方法首先初始化损失为None，然后检查是否提供了closure，如果是，则调用它来计算损失。接下来，对每个参数组进行梯度更新。对于每个参数，首先检查梯度是否为None，如果是，则继续到下一个参数。然后，根据权重衰减因子对梯度进行更新。如果动量不为0，则根据参数的状态更新动量缓冲区，并将动量缓冲区与梯度进行更新。如果启用了Nesterov动量，则将梯度与动量缓冲区的乘积添加到梯度中，否则直接使用动量缓冲区的值。最后，将参数的值减去梯度的乘积。
这个类的作用是实现随机梯度下降算法，可以用于优化模型的参数。它支持带有动量的更新
这段代码是SGD_caffe类中的step方法，用于执行单次优化步骤。
该方法接受一个可调用对象closure，用于重新评估模型并返回损失值。在方法开始时，初始化损失为None。然后，如果提供了closure，则调用它来计算损失。
接下来，对每个参数组进行循环迭代。对于每个参数组，获取相关的参数，包括权重衰减因子weight_decay、动量因子momentum、动量的阻尼因子dampening和是否启用Nesterov动量的标志nesterov。
然后，对于参数组中的每个参数p，首先检查其梯度是否为None，如果是，则继续到下一个参数。然后，根据权重衰减因子对梯度进行更新。如果动量因子不为0，则获取参数的状态param_state，如果参数的状态中不存在动量缓冲区momentum_buffer，则创建一个与参数p形状相同的全零张量作为动量缓冲区，并根据动量因子和学习率对梯度进行更新。如果动量缓冲区已存在，则直接使用它，并根据动量因子和学习率对梯度进行更新。如果启用了Nesterov动量，则将梯度与动量缓冲区的乘积添加到梯度中，否则直接使用动量缓冲区的值。
最后，将参数的值减去梯度的乘积，实现参数的更新。
方法返回损失值。
这个step方法是优化器在每个优化步骤中执行的操作，用于更新模型的参数。它根据参数组的配置进行权重衰减、动量计算和参数更新。
import os, sys
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.models as models
import torch.autograd.variable as Variable
import numpy as np
import scipy.io as sio
from os.path import join as pjoin
#from skimage.transform import resize
#from models import HiFi1Edge
import skimage.io as io
import time
import skimage
import warnings
from PIL import Image
class Logger(object):
def __init__(self, fpath=None):
self.console = sys.stdout
self.file = None
if fpath is not None:
self.file = open(fpath, 'w')
def __del__(self):
self.close()
def __enter__(self):
pass
def __exit__(self, *args):
self.close()
def write(self, msg):
self.console.write(msg)
if self.file is not None:
self.file.write(msg)
def flush(self):
self.console.flush()
if self.file is not None:
self.file.flush()
os.fsync(self.file.fileno())
def close(self):
self.console.close()
if self.file is not None:
self.file.close()
class Averagvalue(object):
"""Computes and stores the average and current value"""
def __init__(self):
self.reset()
def reset(self):
self.val = 0
self.avg = 0
self.sum = 0
self.count = 0
def update(self, val, n=1):
self.val = val
self.sum += val * n
self.count += n
self.avg = self.sum / self.count
def save_checkpoint(state, filename='checkpoint.pth'):
torch.save(state, filename)
def load_pretrained(model, fname, optimizer=None):
"""
resume training from previous checkpoint
:param fname: filename(with path) of checkpoint file
:return: model, optimizer, checkpoint epoch
"""
if os.path.isfile(fname):
print("=> loading checkpoint '{}'".format(fname))
checkpoint = torch.load(fname)
model.load_state_dict(checkpoint['state_dict'])
if optimizer is not None:
optimizer.load_state_dict(checkpoint['optimizer'])
return model, optimizer, checkpoint['epoch']
else:
return model, checkpoint['epoch']
else:
print("=> no checkpoint found at '{}'".format(fname))
def load_vgg16pretrain(model, vggmodel='vgg16convs.mat'):
vgg16 = sio.loadmat(vggmodel)
torch_params = model.state_dict()
for k in vgg16.keys():
name_par = k.split('-')
size = len(name_par)
if size == 2:
name_space = name_par[0] + '.' + name_par[1]
data = np.squeeze(vgg16[k])
torch_params[name_space] = torch.from_numpy(data)
model.load_state_dict(torch_params)
def load_vgg16pretrain_half(model, vggmodel='vgg16convs.mat'):
vgg16 = sio.loadmat(vggmodel)
torch_params = model.state_dict()
for k in vgg16.keys():
name_par = k.split('-')
size = len(name_par)
if size == 2:
name_space = name_par[0] + '.' + name_par[1]
data = np.squeeze(vgg16[k])
shape = data.shape
index = int(shape[0]/2)
if len(shape) == 1:
data = data[:index]
else:
data = data[:index,:,:,:]
torch_params[name_space] = torch.from_numpy(data)
model.load_state_dict(torch_params)
def load_fsds_caffe(model, fsdsmodel='caffe-fsds.mat'):
fsds = sio.loadmat(fsdsmodel)
torch_params = model.state_dict()
for k in fsds.keys():
name_par = k.split('-')
size = len(name_par)
data = np.squeeze(fsds[k])
if 'upsample' in name_par:
# print('skip upsample')
continue
if size == 2:
name_space = name_par[0] + '.' + name_par[1]
data = np.squeeze(fsds[k])
if data.ndim==2:
data = np.reshape(data, (data.shape[0], data.shape[1]))
torch_params[name_space] = torch.from_numpy(data)
if size == 3:
# if 'bias' in name_par:
# continue
name_space = name_par[0] + '_' + name_par[1]+ '.' + name_par[2]
data = np.squeeze(fsds[k])
# print(data.shape)
if data.ndim==2:
# print (data.shape[0])
data = np.reshape(data,(data.shape[0], data.shape[1]))
if data.ndim==1 :
data = np.reshape(data, (1, len(data), 1, 1))
if data.ndim==0:
data = np.reshape(data, (1))
torch_params[name_space] = torch.from_numpy(data)
if size == 4:
# if 'bias' in name_par:
# continue
data = np.squeeze(fsds[k])
name_space = name_par[0] + '_' + name_par[1] + name_par[2] + '.' + name_par[3]
if data.ndim==2:
data = np.reshape(data,(data.shape[0], data.shape[1], 1, 1))
torch_params[name_space] = torch.from_numpy(data)
model.load_state_dict(torch_params)
print('loaded')
这段代码是一个典型的PyTorch脚本，用于导入所需的库和模块。让我逐行解释一下代码的功能：
import os, sys
这两行导入了Python的os和sys模块，用于处理文件和系统相关的操作。
import torch
import torch.nn as nn
import torch.nn.functional as F
这三行导入了PyTorch库中的核心模块，包括张量操作、神经网络模块和函数。
import torchvision.models as models
这一行导入了PyTorch的torchvision.models模块，该模块包含了一些流行的预训练模型，如VGG、ResNet等。
import torch.autograd.variable as Variable
这一行导入了PyTorch的torch.autograd.variable模块，用于创建可自动求导的变量。
import numpy as np
这一行导入了NumPy库，用于处理数组和矩阵的操作。
import scipy.io as sio
这一行导入了SciPy库的scipy.io模块，用于读取和写入MATLAB文件格式。
from os.path import join as pjoin
这一行导入了os.path模块中的join函数，并将其重命名为pjoin，用于拼接文件路径。
import skimage.io as io
import skimage
这两行导入了Scikit-image库，用于图像的读取和处理。
import warnings
这一行导入了Python的warnings模块，用于处理警告信息。
from PIL import Image
这一行导入了Python Imaging Library（PIL）库的Image模块，用于图像的读取和处理。
这段代码还包含了一些注释部分，被注释掉的代码是一些额外的导入语句和模型定义，可能是因为在这段代码的上下文中没有使用到。
这段代码定义了一个名为Logger的类，用于将输出消息同时打印到控制台和文件中。让我逐行解释一下代码的功能：
class Logger(object):
这一行定义了一个名为Logger的类。
def __init__(self, fpath=None):
self.console = sys.stdout
self.file = None
if fpath is not None:
self.file = open(fpath, 'w')
这是Logger类的构造函数。它接受一个参数fpath，用于指定日志文件的路径。在构造函数中，首先将self.console设置为sys.stdout，即标准输出，表示将日志消息打印到控制台。然后，将self.file初始化为None。如果提供了fpath，则通过open函数创建一个以写入模式打开的文件对象，并将其赋值给self.file。
def __del__(self):
self.close()
这是Logger类的析构函数。它在对象销毁时被调用，并调用self.close()方法关闭日志文件。
def __enter__(self):
pass
这是Logger类的上下文管理器的__enter__方法。在使用with语句创建Logger对象时，__enter__方法会被调用。在这个方法中，我们不执行任何操作，所以只是简单地使用pass语句占位。
def __exit__(self, *args):
self.close()
这是Logger类的上下文管理器的__exit__方法。在使用with语句结束时，__exit__方法会被调用。在这个方法中，我们调用self.close()方法关闭日志文件。
def write(self, msg):
self.console.write(msg)
if self.file is not None:
self.file.write(msg)
这是Logger类的write方法。它接受一个参数msg，表示要写入的消息。在这个方法中，首先将消息写入到self.console，即控制台输出。然后，检查self.file是否为None，如果不为None，则将消息写入到日志文件中。
def flush(self):
self.console.flush()
if self.file is not None:
self.file.flush()
os.fsync(self.file.fileno())
这是Logger类的flush方法。它用于刷新输出缓冲区。在这个方法中，首先刷新self.console，然后检查self.file是否为None，如果不为None，则刷新日志文件，并通过os.fsync函数将数据同步到磁盘。
def close(self):
self.console.close()
if self.file is not None:
self.file.close()
这是Logger类的close方法。它用于关闭日志文件和控制台输出。在这个方法中，首先关闭self.console，然后检查self.file是否为None，如果不为None，则关闭日志文件。
这个Logger类可以方便地将输出消息同时记录到控制台和文件中。通过使用with语句创建Logger对象，可以确保在使用完毕后自动关闭日志文件。
这是一个名为load_pretrained的函数，用于从先前的检查点文件中恢复训练。函数的参数和功能如下所示：
def load_pretrained(model, fname, optimizer=None):
"""
resume training from previous checkpoint
:param fname: filename(with path) of checkpoint file
:return: model, optimizer, checkpoint epoch
"""
if os.path.isfile(fname):
print("=> loading checkpoint '{}'".format(fname))
checkpoint = torch.load(fname)
model.load_state_dict(checkpoint['state_dict'])
if optimizer is not None:
optimizer.load_state_dict(checkpoint['optimizer'])
return model, optimizer, checkpoint['epoch']
else:
return model, checkpoint['epoch']
else:
print("=> no checkpoint found at '{}'".format(fname))
函数接受三个参数：
model：要加载预训练权重的模型。fname：检查点文件的文件名（包括路径）。optimizer：可选参数，要加载优化器的情况下的优化器对象。函数的功能如下：
torch.load函数加载检查点文件，将返回的字典存储在checkpoint变量中。model.load_state_dict方法将检查点中的模型权重加载到给定的模型中。optimizer参数，则使用optimizer.load_state_dict方法将检查点中的优化器状态加载到给定的优化器中。如果指定的检查点文件不存在，函数将打印一条消息指示找不到检查点文件。
这个函数的作用是方便地从之前保存的检查点文件中恢复模型的训练状态，包括模型权重、优化器状态和训练轮数。
这是一个名为load_vgg16pretrain的函数，用于从预训练的VGG模型文件中加载权重。函数的参数和功能如下所示：
def load_vgg16pretrain(model, vggmodel='vgg16convs.mat'):
vgg16 = sio.loadmat(vggmodel)
torch_params = model.state_dict()
for k in vgg16.keys():
name_par = k.split('-')
size = len(name_par)
if size == 2:
name_space = name_par[0] + '.' + name_par[1]
data = np.squeeze(vgg16[k])
torch_params[name_space] = torch.from_numpy(data)
model.load_state_dict(torch_params)
函数接受两个参数：
model：要加载权重的模型。vggmodel：VGG模型文件的文件名（包括路径），默认为’vgg16convs.mat’。函数的功能如下：
scipy.io.loadmat函数加载VGG模型文件，将返回的字典存储在vgg16变量中。model.state_dict方法获取模型的状态字典，并将其存储在torch_params变量中。vgg16字典的键。name_space，其中第一个元素作为模型中的命名空间，第二个元素作为模型中的参数名称。vgg16中获取相应键的值，并使用np.squeeze函数去除多余的维度。torch_params中对应命名空间和参数名称的项。model.load_state_dict方法将更新后的模型状态字典加载到给定的模型中。这个函数的作用是根据预训练的VGG模型文件，将对应的权重加载到给定的模型中。它可以用于在使用VGG模型进行迁移学习或初始化模型权重时快速加载预训练的权重。
这是一个名为load_vgg16pretrain_half的函数，与之前的load_vgg16pretrain函数类似，用于从预训练的VGG模型文件中加载权重。函数的参数和功能如下所示：
def load_vgg16pretrain_half(model, vggmodel='vgg16convs.mat'):
vgg16 = sio.loadmat(vggmodel)
torch_params = model.state_dict()
for k in vgg16.keys():
name_par = k.split('-')
size = len(name_par)
if size == 2:
name_space = name_par[0] + '.' + name_par[1]
data = np.squeeze(vgg16[k])
shape = data.shape
index = int(shape[0]/2)
if len(shape) == 1:
data = data[:index]
else:
data = data[:index,:,:,:]
torch_params[name_space] = torch.from_numpy(data)
model.load_state_dict(torch_params)
函数接受两个参数：
model：要加载权重的模型。vggmodel：VGG模型文件的文件名（包括路径），默认为’vgg16convs.mat’。函数的功能如下：
scipy.io.loadmat函数加载VGG模型文件，将返回的字典存储在vgg16变量中。model.state_dict方法获取模型的状态字典，并将其存储在torch_params变量中。vgg16字典的键。name_space，其中第一个元素作为模型中的命名空间，第二个元素作为模型中的参数名称。vgg16中获取相应键的值，并使用np.squeeze函数去除多余的维度。torch_params中对应命名空间和参数名称的项。model.load_state_dict方法将更新后的模型状态字典加载到给定的模型中。这个函数与load_vgg16pretrain函数的不同之处在于，它只加载VGG模型权重的一半，用于模型的快速初始化或迁移学习。
这是一个名为load_vgg16pretrain_half的函数，与之前的load_vgg16pretrain函数类似，用于从预训练的VGG模型文件中加载权重。函数的参数和功能如下所示：
def load_vgg16pretrain_half(model, vggmodel='vgg16convs.mat'):
vgg16 = sio.loadmat(vggmodel)
torch_params = model.state_dict()
for k in vgg16.keys():
name_par = k.split('-')
size = len(name_par)
if size == 2:
name_space = name_par[0] + '.' + name_par[1]
data = np.squeeze(vgg16[k])
shape = data.shape
index = int(shape[0]/2)
if len(shape) == 1:
data = data[:index]
else:
data = data[:index,:,:,:]
torch_params[name_space] = torch.from_numpy(data)
model.load_state_dict(torch_params)
函数接受两个参数：
model：要加载权重的模型。vggmodel：VGG模型文件的文件名（包括路径），默认为’vgg16convs.mat’。函数的功能如下：
scipy.io.loadmat函数加载VGG模型文件，将返回的字典存储在vgg16变量中。model.state_dict方法获取模型的状态字典，并将其存储在torch_params变量中。vgg16字典的键。name_space，其中第一个元素作为模型中的命名空间，第二个元素作为模型中的参数名称。vgg16中获取相应键的值，并使用np.squeeze函数去除多余的维度。torch_params中对应命名空间和参数名称的项。model.load_state_dict方法将更新后的模型状态字典加载到给定的模型中。这个函数与load_vgg16pretrain函数的不同之处在于，它只加载VGG模型权重的一半，用于模型的快速初始化或迁移学习。
这是一个名为load_fsds_caffe的函数，用于从预训练的FSDS Caffe模型文件中加载权重。函数的参数和功能如下所示：
def load_fsds_caffe(model, fsdsmodel='caffe-fsds.mat'):
fsds = sio.loadmat(fsdsmodel)
torch_params = model.state_dict()
for k in fsds.keys():
name_par = k.split('-')
size = len(name_par)
data = np.squeeze(fsds[k])
if 'upsample' in name_par:
continue
if size == 2:
name_space = name_par[0] + '.' + name_par[1]
if data.ndim == 2:
data = np.reshape(data, (data.shape[0], data.shape[1]))
torch_params[name_space] = torch.from_numpy(data)
if size == 3:
name_space = name_par[0] + '_' + name_par[1]+ '.' + name_par[2]
if data.ndim == 2:
data = np.reshape(data, (data.shape[0], data.shape[1]))
if data.ndim == 1:
data = np.reshape(data, (1, len(data), 1, 1))
if data.ndim == 0:
data = np.reshape(data, (1))
torch_params[name_space] = torch.from_numpy(data)
if size == 4:
name_space = name_par[0] + '_' + name_par[1] + name_par[2] + '.' + name_par[3]
if data.ndim == 2:
data = np.reshape(data, (data.shape[0], data.shape[1], 1, 1))
torch_params[name_space] = torch.from_numpy(data)
model.load_state_dict(torch_params)
print('loaded')
函数接受两个参数：
model：要加载权重的模型。fsdsmodel：FSDS Caffe模型文件的文件名（包括路径），默认为’caffe-fsds.mat’。函数的功能如下：
scipy.io.loadmat函数加载FSDS Caffe模型文件，将返回的字典存储在fsds变量中。model.state_dict方法获取模型的状态字典，并将其存储在torch_params变量中。fsds字典的键。name_space，其中第一个元素作为模型中的命名空间，第二个元素作为模型中的参数名称。fsds中获取相应键的值，并使用np.squeeze函数去除多余的维度。torch_params中对应命名空间和参数名称的项。name_space，其中第一个、第二个、第三个元素作为模型中的命名空间的一部分。fsds中获取相应键的值，并使用np.squeeze函数去除多余的维度。torch_params中对应命名空间的项。name_space，其中第一个、第二个、第三个、第四个元素作为模型中的命名空间的一部分。fsds中获取相应键的值，并使用np.squeeze函数去除多余的维度。torch_params中对应命名空间的项。model.load_state_dict方法将更新后的模型状态字典加载到给定的模型中，并打印"loaded"。这个函数的作用是根据预训练的FSDS Caffe模型文件，将对应的权重加载到给定的模型中。它可以用于在使用FSDS Caffe模型进行迁移学习或初始化模型权重时快速加载预训练的权重。
#!/user/bin/python
# coding=utf-8
import os, sys
import numpy as np
from PIL import Image
import cv2
import shutil
import argparse
import time
import datetime
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.optim import lr_scheduler
import torchvision
import torchvision.transforms as transforms
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from data_loader import BSDS_RCFLoader
from models import RCF
from functions import cross_entropy_loss_RCF, SGD_caffe
from torch.utils.data import DataLoader, sampler
from utils import Logger, Averagvalue, save_checkpoint, load_vgg16pretrain
from os.path import join, split, isdir, isfile, splitext, split, abspath, dirname
from torchvision import models
from torchsummary import summary
parser = argparse.ArgumentParser(description='PyTorch Training')
parser.add_argument('--batch_size', default=1, type=int, metavar='BT',
help='batch size')
# =============== optimizer
parser.add_argument('--lr', '--learning_rate', default=1e-6, type=float,
metavar='LR', help='initial learning rate')
parser.add_argument('--momentum', default=0.9, type=float, metavar='M',
help='momentum')
parser.add_argument('--weight_decay', '--wd', default=2e-4, type=float,
metavar='W', help='default weight decay')
parser.add_argument('--stepsize', default=3, type=int,
metavar='SS', help='learning rate step size')
parser.add_argument('--gamma', '--gm', default=0.1, type=float,
help='learning rate decay parameter: Gamma')
parser.add_argument('--maxepoch', default=30, type=int, metavar='N',
help='number of total epochs to run')
parser.add_argument('--itersize', default=10, type=int,
metavar='IS', help='iter size')
# =============== misc
parser.add_argument('--start_epoch', default=0, type=int, metavar='N',
help='manual epoch number (useful on restarts)')
parser.add_argument('--print_freq', '-p', default=1000, type=int,
metavar='N', help='print frequency (default: 50)')
parser.add_argument('--gpu', default='0', type=str,
help='GPU ID')
parser.add_argument('--resume', default='', type=str, metavar='PATH',
help='path to latest checkpoint (default: none)')
parser.add_argument('--tmp', help='tmp folder', default='tmp/RCF')
# ================ dataset
# parser.add_argument('--dataset', help='root folder of dataset', default='data/HED-BSDS_PASCAL')
parser.add_argument('--dataset', help='root folder of dataset', default='/home/itlchennai/Edge_Analysis/pack_1/pack_1/data/HED-BSDS/')
args = parser.parse_args()
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" # see issue #152
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
THIS_DIR = abspath(dirname(__file__))
TMP_DIR = join(THIS_DIR, args.tmp)
if not isdir(TMP_DIR):
os.makedirs(TMP_DIR)
# print('***', args.lr)
def main():
args.cuda = True
# dataset
# train_dataset = BSDS_RCFLoader(root=args.dataset, lst= "train_pair.lst")
# test_dataset = BSDS_RCFLoader(root=args.dataset, lst= "test.lst")
train_dataset = BSDS_RCFLoader(root=args.dataset, split= "train")
test_dataset = BSDS_RCFLoader(root=args.dataset, split= "test")
train_loader = DataLoader(
train_dataset, batch_size=args.batch_size,
num_workers=8, drop_last=True,shuffle=True)
test_loader = DataLoader(
test_dataset, batch_size=args.batch_size,
num_workers=8, drop_last=True,shuffle=False)
with open('/home/itlchennai/Edge_Analysis/pack_1/pack_1/data/HED-BSDS/test.lst', 'r') as f:
test_list = f.readlines()
test_list = [split(i.rstrip())[1] for i in test_list]
assert len(test_list) == len(test_loader), "%d vs %d" % (len(test_list), len(test_loader))
# model
model = RCF()
# print(model)
model.cuda()
model.apply(weights_init)
load_vgg16pretrain(model)
if args.resume:
if isfile(args.resume):
# print("=> loading checkpoint '{}'".format(args.resume))
checkpoint = torch.load(args.resume)
model.load_state_dict(checkpoint['state_dict'])
# print("=> loaded checkpoint '{}'".format(args.resume))
# else:
# print("=> no checkpoint found at '{}'".format(args.resume))
#tune lr
net_parameters_id = {}
net = model
for pname, p in net.named_parameters():
if pname in ['conv1_1.weight','conv1_2.weight',
'conv2_1.weight','conv2_2.weight',
'conv3_1.weight','conv3_2.weight','conv3_3.weight',
'conv4_1.weight','conv4_2.weight','conv4_3.weight']:
# print(pname, 'lr:1 de:1')
if 'conv1-4.weight' not in net_parameters_id:
net_parameters_id['conv1-4.weight'] = []
net_parameters_id['conv1-4.weight'].append(p)
elif pname in ['conv1_1.bias','conv1_2.bias',
'conv2_1.bias','conv2_2.bias',
'conv3_1.bias','conv3_2.bias','conv3_3.bias',
'conv4_1.bias','conv4_2.bias','conv4_3.bias']:
# print(pname, 'lr:2 de:0')
if 'conv1-4.bias' not in net_parameters_id:
net_parameters_id['conv1-4.bias'] = []
net_parameters_id['conv1-4.bias'].append(p)
elif pname in ['conv5_1.weight','conv5_2.weight','conv5_3.weight']:
# print(pname, 'lr:100 de:1')
if 'conv5.weight' not in net_parameters_id:
net_parameters_id['conv5.weight'] = []
net_parameters_id['conv5.weight'].append(p)
elif pname in ['conv5_1.bias','conv5_2.bias','conv5_3.bias'] :
# print(pname, 'lr:200 de:0')
if 'conv5.bias' not in net_parameters_id:
net_parameters_id['conv5.bias'] = []
net_parameters_id['conv5.bias'].append(p)
elif pname in ['conv1_1_down.weight','conv1_2_down.weight',
'conv2_1_down.weight','conv2_2_down.weight',
'conv3_1_down.weight','conv3_2_down.weight','conv3_3_down.weight',
'conv4_1_down.weight','conv4_2_down.weight','conv4_3_down.weight',
'conv5_1_down.weight','conv5_2_down.weight','conv5_3_down.weight']:
# print(pname, 'lr:0.1 de:1')
if 'conv_down_1-5.weight' not in net_parameters_id:
net_parameters_id['conv_down_1-5.weight'] = []
net_parameters_id['conv_down_1-5.weight'].append(p)
elif pname in ['conv1_1_down.bias','conv1_2_down.bias',
'conv2_1_down.bias','conv2_2_down.bias',
'conv3_1_down.bias','conv3_2_down.bias','conv3_3_down.bias',
'conv4_1_down.bias','conv4_2_down.bias','conv4_3_down.bias',
'conv5_1_down.bias','conv5_2_down.bias','conv5_3_down.bias']:
# print(pname, 'lr:0.2 de:0')
if 'conv_down_1-5.bias' not in net_parameters_id:
net_parameters_id['conv_down_1-5.bias'] = []
net_parameters_id['conv_down_1-5.bias'].append(p)
elif pname in ['score_dsn1.weight','score_dsn2.weight','score_dsn3.weight',
'score_dsn4.weight','score_dsn5.weight']:
# print(pname, 'lr:0.01 de:1')
if 'score_dsn_1-5.weight' not in net_parameters_id:
net_parameters_id['score_dsn_1-5.weight'] = []
net_parameters_id['score_dsn_1-5.weight'].append(p)
elif pname in ['score_dsn1.bias','score_dsn2.bias','score_dsn3.bias',
'score_dsn4.bias','score_dsn5.bias']:
# print(pname, 'lr:0.02 de:0')
if 'score_dsn_1-5.bias' not in net_parameters_id:
net_parameters_id['score_dsn_1-5.bias'] = []
net_parameters_id['score_dsn_1-5.bias'].append(p)
elif pname in ['score_final.weight']:
# print(pname, 'lr:0.001 de:1')
if 'score_final.weight' not in net_parameters_id:
net_parameters_id['score_final.weight'] = []
net_parameters_id['score_final.weight'].append(p)
elif pname in ['score_final.bias']:
# print(pname, 'lr:0.002 de:0')
if 'score_final.bias' not in net_parameters_id:
net_parameters_id['score_final.bias'] = []
net_parameters_id['score_final.bias'].append(p)
optimizer = torch.optim.SGD([
{'params': net_parameters_id['conv1-4.weight'] , 'lr': args.lr*1 , 'weight_decay': args.weight_decay},
{'params': net_parameters_id['conv1-4.bias'] , 'lr': args.lr*2 , 'weight_decay': 0.},
{'params': net_parameters_id['conv5.weight'] , 'lr': args.lr*100 , 'weight_decay': args.weight_decay},
{'params': net_parameters_id['conv5.bias'] , 'lr': args.lr*200 , 'weight_decay': 0.},
{'params': net_parameters_id['conv_down_1-5.weight'], 'lr': args.lr*0.1 , 'weight_decay': args.weight_decay},
{'params': net_parameters_id['conv_down_1-5.bias'] , 'lr': args.lr*0.2 , 'weight_decay': 0.},
{'params': net_parameters_id['score_dsn_1-5.weight'], 'lr': args.lr*0.01 , 'weight_decay': args.weight_decay},
{'params': net_parameters_id['score_dsn_1-5.bias'] , 'lr': args.lr*0.02 , 'weight_decay': 0.},
{'params': net_parameters_id['score_final.weight'] , 'lr': args.lr*0.001, 'weight_decay': args.weight_decay},
{'params': net_parameters_id['score_final.bias'] , 'lr': args.lr*0.002, 'weight_decay': 0.},
], lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay)
scheduler = lr_scheduler.StepLR(optimizer, step_size=args.stepsize, gamma=args.gamma)
# optimizer = torch.optim.Adam([
# {'params': net_parameters_id['conv1-4.weight'] , 'lr': args.lr*1 , 'weight_decay': args.weight_decay},
# {'params': net_parameters_id['conv1-4.bias'] , 'lr': args.lr*2 , 'weight_decay': 0.},
# {'params': net_parameters_id['conv5.weight'] , 'lr': args.lr*100 , 'weight_decay': args.weight_decay},
# {'params': net_parameters_id['conv5.bias'] , 'lr': args.lr*200 , 'weight_decay': 0.},
# {'params': net_parameters_id['conv_down_1-5.weight'], 'lr': args.lr*0.1 , 'weight_decay': args.weight_decay},
# {'params': net_parameters_id['conv_down_1-5.bias'] , 'lr': args.lr*0.2 , 'weight_decay': 0.},
# {'params': net_parameters_id['score_dsn_1-5.weight'], 'lr': args.lr*0.01 , 'weight_decay': args.weight_decay},
# {'params': net_parameters_id['score_dsn_1-5.bias'] , 'lr': args.lr*0.02 , 'weight_decay': 0.},
# {'params': net_parameters_id['score_final.weight'] , 'lr': args.lr*0.001, 'weight_decay': args.weight_decay},
# {'params': net_parameters_id['score_final.bias'] , 'lr': args.lr*0.002, 'weight_decay': 0.},
# ], lr=args.lr, betas=(0.9, 0.99), weight_decay=args.weight_decay)
# scheduler = lr_scheduler.StepLR(optimizer, step_size=args.stepsize, gamma=args.gamma)
# log
log = Logger(join(TMP_DIR, '%s-%d-log.txt' %('sgd',args.lr)))
sys.stdout = log
train_loss = []
train_loss_detail = []
for epoch in range(args.start_epoch, args.maxepoch):
if epoch == 0:
# print("Performing initial testing...")
multiscale_test(model, test_loader, epoch=epoch, test_list=test_list,
save_dir = join(TMP_DIR, 'initial-testing-record'))
tr_avg_loss, tr_detail_loss = train(
train_loader, model, optimizer, epoch,
save_dir = join(TMP_DIR, 'epoch-%d-training-record' % epoch))
test(model, test_loader, epoch=epoch, test_list=test_list,
save_dir = join(TMP_DIR, 'epoch-%d-testing-record-view' % epoch))
multiscale_test(model, test_loader, epoch=epoch, test_list=test_list,
save_dir = join(TMP_DIR, 'epoch-%d-testing-record' % epoch))
log.flush() # write log
# Save checkpoint
save_file = os.path.join(TMP_DIR, 'checkpoint_epoch{}.pth'.format(epoch))
save_checkpoint({
'epoch': epoch,
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict()
}, filename=save_file)
scheduler.step() # will adjust learning rate
# save train/val loss/accuracy, save every epoch in case of early stop
train_loss.append(tr_avg_loss)
train_loss_detail += tr_detail_loss
def train(train_loader, model, optimizer, epoch, save_dir):
batch_time = Averagvalue()
data_time = Averagvalue()
losses = Averagvalue()
# switch to train mode
model.train()
end = time.time()
epoch_loss = []
counter = 0
for i, (image, label) in enumerate(train_loader):
# measure data loading time
data_time.update(time.time() - end)
image, label = image.cuda(), label.cuda()
outputs = model(image)
loss = torch.zeros(1).cuda()
for o in outputs:
loss = loss + cross_entropy_loss_RCF(o, label)
counter += 1
loss = loss / args.itersize
loss.backward()
if counter == args.itersize:
optimizer.step()
optimizer.zero_grad()
counter = 0
# measure accuracy and record loss
losses.update(loss.item(), image.size(0))
epoch_loss.append(loss.item())
batch_time.update(time.time() - end)
end = time.time()
# display and logging
if not isdir(save_dir):
os.makedirs(save_dir)
if i % args.print_freq == 0:
info = 'Epoch: [{0}/{1}][{2}/{3}] '.format(epoch, args.maxepoch, i, len(train_loader)) + \
'Time {batch_time.val:.3f} (avg:{batch_time.avg:.3f}) '.format(batch_time=batch_time) + \
'Loss {loss.val:f} (avg:{loss.avg:f}) '.format(
loss=losses)
print(info)
label_out = torch.eq(label, 1).float()
outputs.append(label_out)
_, _, H, W = outputs[0].shape
all_results = torch.zeros((len(outputs), 1, H, W))
for j in range(len(outputs)):
all_results[j, 0, :, :] = outputs[j][0, 0, :, :]
torchvision.utils.save_image(1-all_results, join(save_dir, "iter-%d.jpg" % i))
# save checkpoint
save_checkpoint({
'epoch': epoch,
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict()
}, filename=join(save_dir, "epoch-%d-checkpoint.pth" % epoch))
return losses.avg, epoch_loss
def test(model, test_loader, epoch, test_list, save_dir):
model.eval()
if not isdir(save_dir):
os.makedirs(save_dir)
for idx, image in enumerate(test_loader):
image = image.cuda()
_, _, H, W = image.shape
results = model(image)
result = torch.squeeze(results[-1].detach()).cpu().numpy()
results_all = torch.zeros((len(results), 1, H, W))
for i in range(len(results)):
results_all[i, 0, :, :] = results[i]
filename = splitext(test_list[idx])[0]
torchvision.utils.save_image(1-results_all, join(save_dir, "%s.jpg" % filename))
result = Image.fromarray((result * 255).astype(np.uint8))
result.save(join(save_dir, "%s.png" % filename))
# print("Running test [%d/%d]" % (idx + 1, len(test_loader)))
# torch.nn.functional.upsample(input, size=None, scale_factor=None, mode='nearest', align_corners=None)
def multiscale_test(model, test_loader, epoch, test_list, save_dir):
model.eval()
if not isdir(save_dir):
os.makedirs(save_dir)
scale = [0.5, 1, 1.5]
for idx, image in enumerate(test_loader):
image = image[0]
image_in = image.numpy().transpose((1,2,0))
_, H, W = image.shape
multi_fuse = np.zeros((H, W), np.float32)
for k in range(0, len(scale)):
im_ = cv2.resize(image_in, None, fx=scale[k], fy=scale[k], interpolation=cv2.INTER_LINEAR)
im_ = im_.transpose((2,0,1))
results = model(torch.unsqueeze(torch.from_numpy(im_).cuda(), 0))
result = torch.squeeze(results[-1].detach()).cpu().numpy()
fuse = cv2.resize(result, (W, H), interpolation=cv2.INTER_LINEAR)
multi_fuse += fuse
multi_fuse = multi_fuse / len(scale)
### rescale trick suggested by jiangjiang
# multi_fuse = (multi_fuse - multi_fuse.min()) / (multi_fuse.max() - multi_fuse.min())
filename = splitext(test_list[idx])[0]
result_out = Image.fromarray(((1-multi_fuse) * 255).astype(np.uint8))
result_out.save(join(save_dir, "%s.jpg" % filename))
result_out_test = Image.fromarray((multi_fuse * 255).astype(np.uint8))
result_out_test.save(join(save_dir, "%s.png" % filename))
# print("Running test [%d/%d]" % (idx + 1, len(test_loader)))
def weights_init(m):
if isinstance(m, nn.Conv2d):
# xavier(m.weight.data)
m.weight.data.normal_(0, 0.01)
if m.weight.data.shape == torch.Size([1, 5, 1, 1]):
# for new_score_weight
torch.nn.init.constant_(m.weight, 0.2) # as per https://github.com/yun-liu/rcf
if m.bias is not None:
m.bias.data.zero_()
if __name__ == '__main__':
startinitial_time = time.time()
main()
# print("--- %s seconds ---" % (time.time() - startinitial_time))
这段代码导入了一些常用的Python库和自定义的模块，并定义了一些变量和函数。让我们逐行解释每个导入和定义的部分：
import time
import datetime
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.optim import lr_scheduler
import torchvision
import torchvision.transforms as transforms
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
这部分导入了一些常用的Python库，包括time（用于处理时间和日期）、torch（PyTorch深度学习库）、torch.nn（PyTorch中的神经网络模块）、torch.nn.functional（包含各种神经网络函数）、torch.optim.lr_scheduler（PyTorch中的学习率调度器模块）、torchvision（PyTorch中的计算机视觉库）、torchvision.transforms（用于数据转换的模块，如图像预处理）、matplotlib（用于绘图和可视化的库）。
from data_loader import BSDS_RCFLoader
from models import RCF
from functions import cross_entropy_loss_RCF, SGD_caffe
from torch.utils.data import DataLoader, sampler
from utils import Logger, Averagvalue, save_checkpoint, load_vgg16pretrain
这部分导入了一些自定义的模块和函数。data_loader模块中导入了BSDS_RCFLoader类，用于加载数据集。models模块中导入了RCF类，表示一个自定义的神经网络模型。functions模块中导入了cross_entropy_loss_RCF和SGD_caffe函数，用于定义损失函数和优化器。torch.utils.data模块中导入了DataLoader和sampler类，用于数据加载和采样。utils模块中导入了Logger、Averagvalue、save_checkpoint和load_vgg16pretrain函数，用于日志记录、计算平均值、保存和加载模型等实用功能。
from os.path import join, split, isdir, isfile, splitext, split, abspath, dirname
from torchvision import models
from torchsummary import summary
这部分导入了一些用于文件路径操作的函数和模块，如os.path中的函数。torchvision.models模块导入了一些流行的预训练模型，如VGG、ResNet等。torchsummary模块导入了summary函数，用于打印模型的摘要信息。
以上是代码中的导入语句，用于引入所需的库、模块和函数。接下来的代码应该是使用这些导入的库和模块进行具体的操作，如定义模型、加载数据、训练模型等。如果您有进一步的问题或需要更详细的解释，请随时提问。
这段代码使用了argparse模块来解析命令行参数，用于控制PyTorch训练的一些设置。让我们逐行解释每个部分：
import argparse
parser = argparse.ArgumentParser(description='PyTorch Training')
首先，导入了argparse模块，并创建了一个ArgumentParser对象，命名为parser。ArgumentParser用于解析命令行参数并生成帮助信息。
parser.add_argument('--batch_size', default=1, type=int, metavar='BT',
help='mini-batch size (default: 1)')
接下来，使用add_argument方法添加了一个命令行参数--batch_size。这个参数用于设置每个小批量（mini-batch）的大小。参数的具体设置如下：
--batch_size：参数的名称，前面加上了--表示它是一个命令行参数。default=1：参数的默认值为1。type=int：参数的类型为整数。metavar='BT'：在帮助信息中显示的参数名称。help='mini-batch size (default: 1)'：参数的帮助信息，描述了参数的作用和默认值。通过添加add_argument语句，可以添加更多的命令行参数，用于控制训练过程中的其他设置，如学习率、迭代次数等。
args = parser.parse_args()
最后，使用parse_args方法解析命令行参数，并将解析结果存储在args变量中。args是一个命名空间（Namespace）对象，可以通过属性访问命令行参数的值。
通过使用argparse模块，可以在命令行中指定训练过程的各种参数，而不需要直接修改代码。这样可以提高代码的灵活性和可重用性。
这段代码添加了更多的命令行参数，用于控制PyTorch训练的各种设置。让我们逐行解释每个部分：
parser.add_argument('--lr', '--learning_rate', default=1e-6, type=float,
metavar='LR', help='initial learning rate')
这行代码添加了一个命令行参数--lr（或--learning_rate），用于设置初始学习率。参数的具体设置如下：
default=1e-6：参数的默认值为1e-6。type=float：参数的类型为浮点数。metavar='LR'：在帮助信息中显示的参数名称。help='initial learning rate'：参数的帮助信息，描述了参数的作用和默认值。类似地，下面的代码添加了其他的命令行参数，用于控制训练过程中的不同设置：
parser.add_argument('--momentum', default=0.9, type=float, metavar='M',
help='momentum')
parser.add_argument('--weight_decay', '--wd', default=2e-4, type=float,
metavar='W', help='default weight decay')
parser.add_argument('--stepsize', default=3, type=int,
metavar='SS', help='learning rate step size')
parser.add_argument('--gamma', '--gm', default=0.1, type=float,
help='learning rate decay parameter: Gamma')
parser.add_argument('--maxepoch', default=30, type=int, metavar='N',
help='number of total epochs to run')
parser.add_argument('--itersize', default=10, type=int,
metavar='IS', help='iter size')
这些参数分别用于设置动量（momentum）、权重衰减（weight_decay）、学习率步长（stepsize）、学习率衰减参数（gamma）、总的训练轮数（maxepoch）和迭代大小（itersize）。每个参数都有默认值、类型和帮助信息。
args = parser.parse_args()
最后，使用parse_args方法解析命令行参数，并将解析结果存储在args变量中。args是一个命名空间（Namespace）对象，可以通过属性访问命令行参数的值。
通过添加这些命令行参数，可以在命令行中指定训练过程的各种设置，例如学习率、动量、权重衰减等。这样可以方便地调整训练过程的参数，而不需要直接修改代码。
这段代码添加了更多的命令行参数，用于控制PyTorch训练的其他设置。让我们逐行解释每个部分：
parser.add_argument('--start_epoch', default=0, type=int, metavar='N',
help='manual epoch number (useful on restarts)')
这行代码添加了一个命令行参数--start_epoch，用于设置手动指定的起始轮数。这个参数在重新启动训练时非常有用。参数的具体设置如下：
default=0：参数的默认值为0。type=int：参数的类型为整数。metavar='N'：在帮助信息中显示的参数名称。help='manual epoch number (useful on restarts)'：参数的帮助信息，描述了参数的作用和默认值。类似地，下面的代码添加了其他的命令行参数，用于控制训练过程中的不同设置：
parser.add_argument('--print_freq', '-p', default=1000, type=int,
metavar='N', help='print frequency (default: 50)')
parser.add_argument('--gpu', default='0', type=str,
help='GPU ID')
parser.add_argument('--resume', default='', type=str, metavar='PATH',
help='path to latest checkpoint (default: none)')
parser.add_argument('--tmp', help='tmp folder', default='tmp/RCF')
parser.add_argument('--dataset', help='root folder of dataset', default='/home/itlchennai/Edge_Analysis/pack_1/pack_1/data/HED-BSDS/')
这些参数分别用于设置打印频率（print_freq）、GPU ID（gpu）、恢复训练时的检查点路径（resume）、临时文件夹路径（tmp）和数据集根文件夹路径（dataset）。每个参数都有默认值、类型和帮助信息。
args = parser.parse_args()
最后，使用parse_args方法解析命令行参数，并将解析结果存储在args变量中。args是一个命名空间（Namespace）对象，可以通过属性访问命令行参数的值。
通过添加这些命令行参数，可以在命令行中指定训练过程的各种设置，例如起始轮数、打印频率、GPU ID等。这样可以方便地调整训练过程的参数，而不需要直接修改代码。
这段代码设置了环境变量，用于指定使用的GPU设备。让我们逐行解释每个部分：
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" # see issue #152
这行代码设置了环境变量CUDA_DEVICE_ORDER，用于指定GPU设备的顺序。PCI_BUS_ID表示按照PCI总线ID的顺序来指定GPU设备。这样可以确保使用的GPU设备与命令行参数args.gpu中指定的GPU ID对应。
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
这行代码设置了环境变量CUDA_VISIBLE_DEVICES，用于指定可见的GPU设备。args.gpu是命令行参数中指定的GPU ID，它可以是一个数字（例如0）或多个数字的字符串（例如"0,1,2"）。通过设置CUDA_VISIBLE_DEVICES，可以限制程序使用的GPU设备，只有指定的设备会被程序看到和使用。
通过设置这些环境变量，可以确保程序在训练过程中使用正确的GPU设备。这对于多GPU环境下的训练很有用，可以根据需要灵活地选择使用哪些GPU设备。
这段代码用于创建临时文件夹，并确保临时文件夹存在。让我们逐行解释每个部分：
THIS_DIR = abspath(dirname(__file__))
这行代码获取当前脚本文件的绝对路径所在的目录。abspath函数用于获取绝对路径，dirname函数用于获取目录路径。
TMP_DIR = join(THIS_DIR, args.tmp)
这行代码将临时文件夹路径设置为当前脚本文件所在目录与命令行参数args.tmp拼接而成。join函数用于拼接路径。
if not isdir(TMP_DIR):
os.makedirs(TMP_DIR)
这段代码检查临时文件夹是否存在，如果不存在则创建它。isdir函数用于判断路径是否为一个目录，os.makedirs函数用于递归创建目录。
通过这段代码，可以确保临时文件夹存在并可用于存储临时文件或中间结果。这在训练过程中可能会使用到，例如保存训练过程中的检查点、日志文件等。
这段代码使用argparse模块添加了一个命令行参数--dataset，用于指定数据集的根文件夹路径。让我们逐个解释每个部分：
parser.add_argument('--dataset', help='root folder of dataset', default='/home/itlchennai/Edge_Analysis/pack_1/pack_1/data/HED-BSDS/')
这行代码使用add_argument方法添加了一个命令行参数--dataset。--dataset是参数的名称，可以在命令行中使用该名称来指定参数。help参数用于提供关于该参数的描述信息，当用户在命令行中使用--help参数时，会显示这些描述信息。default参数用于设置参数的默认值，如果用户没有在命令行中指定该参数，就会使用默认值。
在这个例子中，--dataset参数用于指定数据集的根文件夹路径。用户可以在命令行中使用--dataset参数并提供一个路径来指定数据集的根文件夹。如果用户没有指定该参数，就会使用默认的路径/home/itlchennai/Edge_Analysis/pack_1/pack_1/data/HED-BSDS/作为数据集的根文件夹路径。
通过使用argparse模块，可以方便地从命令行中解析参数，并在代码中使用这些参数来控制程序的行为。在这个例子中，--dataset参数可以用于指定不同的数据集路径，从而灵活地处理不同的数据集。
这段代码定义了一个名为main的函数，用于设置和加载数据集。让我们逐行解释每个部分：
def main():
args.cuda = True
这行代码将args.cuda设置为True，表示使用GPU加速。这个参数可能在代码的其他部分用于控制是否使用GPU进行训练。
train_dataset = BSDS_RCFLoader(root=args.dataset, split= "train")
test_dataset = BSDS_RCFLoader(root=args.dataset, split= "test")
这两行代码分别创建了训练集和测试集的数据集对象。BSDS_RCFLoader是一个自定义的数据集加载器，根据给定的参数root和split来加载相应的数据集。root参数指定数据集的根文件夹路径，split参数指定数据集的划分（例如训练集或测试集）。
train_loader = DataLoader(
train_dataset, batch_size=args.batch_size,
num_workers=8, drop_last=True, shuffle=True)
test_loader = DataLoader(
test_dataset, batch_size=args.batch_size,
num_workers=8, drop_last=True, shuffle=False)
这两段代码分别创建了训练集和测试集的数据加载器。DataLoader是PyTorch提供的用于加载数据的工具，它可以自动进行数据批处理、多线程加载等操作。这里使用DataLoader来加载训练集和测试集的数据，其中指定了批处理大小(batch_size)、线程数量(num_workers)、是否舍弃最后一个不完整的批次(drop_last)以及是否对数据进行洗牌(shuffle)。
with open('/home/itlchennai/Edge_Analysis/pack_1/pack_1/data/HED-BSDS/test.lst', 'r') as f:
test_list = f.readlines()
test_list = [split(i.rstrip())[1] for i in test_list]
assert len(test_list) == len(test_loader), "%d vs %d" % (len(test_list), len(test_loader`)`)
这段代码打开一个文件test.lst，读取其中的内容，并将每一行的内容添加到test_list中。然后，对test_list中的每个元素应用split函数，并取返回结果的第二个元素，将结果存储回test_list。最后，代码使用断言(assert)检查test_list的长度是否与test_loader的长度相等。
这段代码的目的可能是为了获取测试数据集的文件列表，并进行一些验证操作，以确保数据集加载正确。
这些代码片段一起构成了main函数的主要功能，用于设置和加载数据集，为后续的训练和测试准备数据。