本周主要学习图像描述,即为图片生成描述语言,输入某一张图片,输出的是一些客观描述图片内容的句子,他可以理解为一种特殊的机器翻译。模型需要有复杂的场景理解能力。但是这种是非常有挑战的,图片是捕捉的真实世界的原始刻画,而自然语言是代表更高一级的抽象。运行一个U-NET模型训练。阅读一篇U-Net论文,用于生物医学影像分割的卷积网络。毕设方面,学弟已经掌握区块链网络的搭建,并且开始搭建前端网页。
普遍认为深度网络的成功训练需要数千个标记好的训练样本。在本文中,我们提出了一种网络和训练策略,依靠高效的数据增强,以更有效地利用现有的标记样本。该结构由捕获上下文的收缩路径和对称的支持精确定位的展开路径组成。我们证明这样的网络可以从很少的图像中进行端到端的训练,并在ISBI电镜神经元结构分割挑战中优于之前最好的方法(滑动窗口卷积网络)。用该网络在透射光镜图像(相衬和DIC)上进行训练,我们在ISBI2015细胞追踪挑战赛中以巨大优势胜出。此外,网络性能也很高。在最新的GPU上,512x512图像的分割所需时间不到一秒
所以该文为了解决医学图像分割问题;提出了一种数据增强方法来有效利用标注数据;提出了一种U型的网络结构可以同时获取上下文信息和位置信息。
在过去的两年里,深度卷积网络在许多视觉识别任务中表现都优于SOTA,例如[7,3]。虽然卷积网络[8]已经存在很长时间了,但是考虑可用训练集的大小和网络的大小,其成功程度依然有限。Krizhevsky等人[7]的突破在于对拥有100万张训练图像的ImageNet数据集进行8层数百万个参数的大型网络的有监督训练。从那时起,出现了更多更大更深的网络[12]。
卷积网络的典型应用是在分类任务上,其中图像的输出是一个类标签。然而,在许多视觉任务中,特别是在生物医学影像处理中,期望输出应包括定位,也就是说,应该给每个像素指定一个类标签。此外,在生物医学任务中,通常也难以得到数以千计的训练图像。因此,Ciresan等人[1]在滑动窗口中训练网络,通过以像素周围的局部区域(patch)作为输入,来预测每个像素的类标签。首先,这个网络可以本地化。其次,在patch方面的训练数据远多于训练图像的数量。该网络在ISBI 2012上以巨大的优势赢得了EM分割挑战。
不过,显然Ciresan等人的[1]策略有两个缺点。首先,它的速度非常慢,因为每个patch都必须单独运行网络,而由于patch有重叠,存在大量的冗余现象。其次,在定位精度和上下文使用之间需要权衡。较大的patch需要更多的max-pooling层,导致定位精度降低;而较小的patch只允许网络看到少量上下文。最近的方法[11,4]提出了一种利用多层特征的分类器输出,可以同时做到良好的本地化与上下文使用。
在本文中,我们构建了一种更优雅的结构,即所谓的“全卷积网络”[9]。我们修改并扩展了这个结构,使其能在少量训练图像下工作,并产生更精确的分割,见图1;[9]中的主要思想是在连续的层中补充一般的收缩网络,在该层中,池化运算符被上采样运算符代替。因此,这些层提高了输出的分辨率。将收缩路径的高分辨率特征与上采样输出相结合,实现了本地化。连续的卷积层可以根据这些信息组合出更准确的输出。
如图:用DIC(微分干涉对比)显微镜记录的玻璃上的HeLa细胞。( a )原始图像( b )覆盖与ground truth分割。不同的颜色表示不同的HeLa细胞。( c )生成的分割蒙版(白色:前景,黑色:背景)。( d )使用逐像素损失权重进行映射,以强制网络学习边缘像素。
所得到的网络适用于各种生物医学分割问题。在本文中,我们展示了EM栈中神经元结构分割的结果(一项从ISBI2012开始至今的比赛),在那里我们超过了Ciresan等人[1]的网络。此外,我们展示了ISBI2015细胞追踪比赛的光镜图像细胞分割结果。在这两个最具挑战性的2D透射光数据集上,我们以较大的优势胜出。
总结:
医学领域图像分割标注数据相对不足;
本文的比较对象为Ciresan et al. [1],该文章通过输入以某个像素点为中心的一个patch以获得该像素点的label,但存在两点不足:1)由于需要逐patch地输入来进行预测,因此非常的慢;2)没有解决位置信息和上下文信息之间的trade-off问题,即大patch有上下文信息但是缺少位置信息(max-pooling所致),小patch有位置信息但是缺少上下文信息;
本文的方法基于FCN[2]。
采用了Overlap-tile strategy:
即由于边界区域的像素缺乏上下文信息,通过在原图像外围“tile”一圈的做法来补全上下文,举例来说,譬如要补全上图中黄框区域的上下文成蓝框区域,具体的做法是将黄框和蓝框之间右侧和下侧的像素通过镜像拷贝的方式拷贝到左侧和上侧,以补全蓝框。
数据增强策略:通过对原始图像进行弹性形变以获得补充图像,这可以让网络学习弹性形变不变性;
加权Loss:增大对粘连的同类物体之间的“background”像素的loss权重,使得每个物体的分割轮廓是清晰的。
Encoder:左半部分,由两个3x3的卷积层(ReLU)+2x2的max polling层(stride=2)反复组成,每经过一次下采样,通道数翻倍;
Decoder:右半部分,由一个2x2的上采样卷积层(ReLU)+Concatenation(crop[3]对应的Encoder层的输出feature map然后与Decoder层的上采样结果相加)+2个3x3的卷积层(ReLU)反复构成;
最后一层通过一个1x1卷积将通道数变成期望的类别数。
在不同的生物医学分割应用中,u-net架构都取得了很好的性能。由于使用了带弹性形变的数据增强,它只需要少量标注好的图像,就能在NVidia Titan(6gb)上达到10小时这一合理的训练时间。我们提供了完整的基于Caffe的实现与训练好的网络。我们相信u-net架构可以很容易地应用到更多的任务中。
本文也是分割领域很经典的一篇paper,UNet基于FCN,对FCN的基本结构进行了更精细的设计,更为高效,是可以替代FCN的方案;
本文采用的Overlap-tile策略、数据增强策略、加权Loss策略等都是非常经典的trick,值得初学者学习借鉴。
import os
import time
import datetime
import torch
from src import UNet
from train_utils import train_one_epoch, evaluate, create_lr_scheduler
from my_dataset import DriveDataset
import transforms as T
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
class SegmentationPresetTrain:
def __init__(self, base_size, crop_size, hflip_prob=0.5, vflip_prob=0.5,
mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)):
min_size = int(0.5 * base_size)
max_size = int(1.2 * base_size)
trans = [T.RandomResize(min_size, max_size)]
if hflip_prob > 0:
trans.append(T.RandomHorizontalFlip(hflip_prob))
if vflip_prob > 0:
trans.append(T.RandomVerticalFlip(vflip_prob))
trans.extend([
T.RandomCrop(crop_size),
T.ToTensor(),
T.Normalize(mean=mean, std=std),
])
self.transforms = T.Compose(trans)
def __call__(self, img, target):
return self.transforms(img, target)
class SegmentationPresetEval:
def __init__(self, mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)):
self.transforms = T.Compose([
T.ToTensor(),
T.Normalize(mean=mean, std=std),
])
def __call__(self, img, target):
return self.transforms(img, target)
def get_transform(train, mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)):
base_size = 565
crop_size = 480
if train:
return SegmentationPresetTrain(base_size, crop_size, mean=mean, std=std)
else:
return SegmentationPresetEval(mean=mean, std=std)
def create_model(num_classes):
model = UNet(in_channels=3, num_classes=num_classes, base_c=32)
return model
def main(args):
device = torch.device(args.device if torch.cuda.is_available() else "cpu")
device = torch.device("cpu")
batch_size = args.batch_size
# segmentation nun_classes + background
num_classes = args.num_classes + 1
# using compute_mean_std.py
mean = (0.709, 0.381, 0.224)
std = (0.127, 0.079, 0.043)
# 用来保存训练以及验证过程中信息
results_file = "results{}.txt".format(datetime.datetime.now().strftime("%Y%m%d-%H%M%S"))
train_dataset = DriveDataset(args.data_path,
train=True,
transforms=get_transform(train=True, mean=mean, std=std))
val_dataset = DriveDataset(args.data_path,
train=False,
transforms=get_transform(train=False, mean=mean, std=std))
num_workers = min([os.cpu_count(), batch_size if batch_size > 1 else 0, 8])
num_workers = 0
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=batch_size,
num_workers=num_workers,
shuffle=True,
pin_memory=True,
collate_fn=train_dataset.collate_fn)
val_loader = torch.utils.data.DataLoader(val_dataset,
batch_size=1,
num_workers=num_workers,
pin_memory=True,
collate_fn=val_dataset.collate_fn)
model = create_model(num_classes=num_classes)
model.to(device)
params_to_optimize = [p for p in model.parameters() if p.requires_grad]
optimizer = torch.optim.SGD(
params_to_optimize,
lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay
)
scaler = torch.cuda.amp.GradScaler() if args.amp else None
# 创建学习率更新策略,这里是每个step更新一次(不是每个epoch)
lr_scheduler = create_lr_scheduler(optimizer, len(train_loader), args.epochs, warmup=True)
if args.resume:
checkpoint = torch.load(args.resume, map_location='cpu')
model.load_state_dict(checkpoint['model'])
optimizer.load_state_dict(checkpoint['optimizer'])
lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
args.start_epoch = checkpoint['epoch'] + 1
if args.amp:
scaler.load_state_dict(checkpoint["scaler"])
best_dice = 0.
start_time = time.time()
for epoch in range(args.start_epoch, args.epochs):
mean_loss, lr = train_one_epoch(model, optimizer, train_loader, device, epoch, num_classes,
lr_scheduler=lr_scheduler, print_freq=args.print_freq, scaler=scaler)
confmat, dice = evaluate(model, val_loader, device=device, num_classes=num_classes)
val_info = str(confmat)
print(val_info)
print(f"dice coefficient: {dice:.3f}")
# write into txt
with open(results_file, "a") as f:
# 记录每个epoch对应的train_loss、lr以及验证集各指标
train_info = f"[epoch: {epoch}]\n" \
f"train_loss: {mean_loss:.4f}\n" \
f"lr: {lr:.6f}\n" \
f"dice coefficient: {dice:.3f}\n"
f.write(train_info + val_info + "\n\n")
if args.save_best is True:
if best_dice < dice:
best_dice = dice
else:
continue
save_file = {"model": model.state_dict(),
"optimizer": optimizer.state_dict(),
"lr_scheduler": lr_scheduler.state_dict(),
"epoch": epoch,
"args": args}
if args.amp:
save_file["scaler"] = scaler.state_dict()
if args.save_best is True:
torch.save(save_file, "save_weights/best_model.pth")
else:
torch.save(save_file, "save_weights/model_{}.pth".format(epoch))
total_time = time.time() - start_time
total_time_str = str(datetime.timedelta(seconds=int(total_time)))
print("training time {}".format(total_time_str))
def parse_args():
import argparse
parser = argparse.ArgumentParser(description="pytorch unet training")
parser.add_argument("--data-path", default="C:\\Users\\666\\Desktop", help="DRIVE root")
# exclude background
parser.add_argument("--num-classes", default=1, type=int)
parser.add_argument("--device", default="cuda", help="training device")
parser.add_argument("-b", "--batch-size", default=4, type=int)
parser.add_argument("--epochs", default=200, type=int, metavar="N",
help="number of total epochs to train")
parser.add_argument('--lr', default=0.01, type=float, help='initial learning rate')
parser.add_argument('--momentum', default=0.9, type=float, metavar='M',
help='momentum')
parser.add_argument('--wd', '--weight-decay', default=1e-4, type=float,
metavar='W', help='weight decay (default: 1e-4)',
dest='weight_decay')
parser.add_argument('--print-freq', default=1, type=int, help='print frequency')
parser.add_argument('--resume', default='', help='resume from checkpoint')
parser.add_argument('--start-epoch', default=0, type=int, metavar='N',
help='start epoch')
parser.add_argument('--save-best', default=True, type=bool, help='only save best dice weights')
# Mixed precision training parameters
parser.add_argument("--amp", default=False, type=bool,
help="Use torch.cuda.amp for mixed precision training")
args = parser.parse_args()
return args
if __name__ == '__main__':
args = parse_args()
if not os.path.exists("./save_weights"):
os.mkdir("./save_weights")
main(args)