我自己的原文哦~ https://blog.51cto.com/whaosoft/11878447
一、PyTorch与torch-xla的桥接
文章从XLATensor开始的溯源、注册PyTorch库实现、从PyTorch调用到torch_xla三个方面来介绍PyTorch与torch-xla的桥接
XLA (Accelerated Linear Algebra)是一个开源的机器学习编译器，对PyTorch、Tensorflow、JAX等多个深度学习框架都有支持。最初XLA实际上是跟Tensorflow深度结合的，很好地服务了Tensorflow和TPU，而与XLA的结合主要依赖于社区的支持，即torch-xla。
torch-xla在支持XLA编译的基础上，较大限度地保持了PyTorch的易用性，贴一个官方的DDP训练的例子：
import torch.distributed as dist
-import torch.multiprocessing as mp
+import torch_xla.core.xla_model as xm
+import torch_xla.distributed.parallel_loader as pl
+import torch_xla.distributed.xla_multiprocessing as xmp
+import torch_xla.distributed.xla_backend def _mp_fn(rank, world_size): ... - os.environ['MASTER_ADDR'] = 'localhost'
- os.environ['MASTER_PORT'] = '12355'
- dist.init_process_group("gloo", rank=rank, world_size=world_size)
+ # Rank and world size are inferred from the XLA device runtime
+ dist.init_process_group("xla", init_method='xla://')
+
+ model.to(xm.xla_device())
+ # `gradient_as_bucket_view=True` required for XLA
+ ddp_model = DDP(model, gradient_as_bucket_view=True) - model = model.to(rank)
- ddp_model = DDP(model, device_ids=[rank])
+ xla_train_loader = pl.MpDeviceLoader(train_loader, xm.xla_device()) - for inputs, labels in train_loader:
+ for inputs, labels in xla_train_loader: optimizer.zero_grad() outputs = ddp_model(inputs) loss = loss_fn(outputs, labels) loss.backward() optimizer.step() if __name__ == '__main__':
- mp.spawn(_mp_fn, args=(), nprocs=world_size)
+ xmp.spawn(_mp_fn, args=())将一段PyTorch代码改写为torch-xla代码，主要就是三个方面：
- 将模型和数据放到xla device上
- 适当的时候调用
xm.mark_step - 某些组件该用pytorchx-xla提供的，比如amp和spawn
其中第二条并没有在上面的代码中体现，原因是为了让用户少改代码，torch-xla将mark_step封装到了dataloader中，实际上不考虑DDP的完整训练的过程可以简写如下：
device = xm.xla_device()
model = model.to(device)
for data, label in enumerate(dataloader): data, label = data.to(device), label.to(device) output = model(data) loss = func(output, label) loss.backward() optimizer.step() xm.mark_step()xm.mark_step的作用就是"告诉"框架：现在对图的定义告一段落了，可以编译并执行计算了。既然如此，那么mark_step之前的内容是做了什么呢？因为要在mark_step之后才编译并计算，那么前面肯定不能执行实际的运算。这就引出了Trace和LazyTensor的概念。
其实到了这里，如果对tensorflow或者torch.fx等比较熟悉，就已经很容易理解了，在mark_step之前，torch-xla将torch Tensor换成了LazyTensor，进而将原本是PyTorch中eager computation的过程替换成了trace的过程，最后生成一张计算图来优化和执行。简而言之这个过程是PyTorch Tensor -> XLATensor -> HLO IR，其中HLO就是XLA所使用的IR。在每次调用到torch op的时候，会调用一次GetIrValue，这时候就意味着一个节点被写入了图中。更具体的信息可以参考XLA Tensor Deep Dive这部分文档。需要注意的是，trace这个过程是独立于mark_step的，即便你的每个循环都不写mark_step，这个循环也可以一直持续下去，只不过在这种情况下，永远都不会发生图的编译和执行，除非在某一步trace的时候，发现图的大小已经超出了pytorch-xla允许的上限。
PyTorch与torch-xla的桥接
知晓了Trace过程之后，就会好奇一个问题：当用户执行一个PyTorch函数调用的时候，torch-xla怎么将这个函数记录下来的？
最容易想到的答案是“torch-xla作为PyTorch的一个编译选项，打开的时候就会使得二者建立起映射关系”，但很可惜，这个答案是错误的，仔细看PyTorch的CMake文件以及torch-xla的编译方式就会明白，torch-xla是几乎单向依赖于PyTorch的（为什么不是全部后面会讲）。既然PyTorch本身在编译期间并不知道torch-xla的存在，那么当用户使用一个xla device上的Tensor作为一个torch function的输入的时候，又经历了怎样一个过程调用到pytorch-xla中的东西呢？
从XLATensor开始的溯源
尽管我们现在并不知道怎么调用到torch-xla中的，但我们知道PyTorch Tensor一定要转换成XLATensor（参考tensor.h），那么我们只需要在关键的转换之处打印出调用堆栈，自然就可以找到调用方，这样虽然不能保证找到PyTorch中的位置，但是能够找到torch-xla中最上层的调用。注意到XLATensor只有下面这一个创建函数接受at::Tensor作为输入，因此就在这里面打印调用栈。
XLATensor XLATensor::Create(const at::Tensor& tensor, const Device& device)测试的用例很简单，我们让两个xla device上的Tensor相乘：
import torch_xla.core.xla_model as xm
import torch device = xm.xla_device()
a = torch.normal(0, 1, (2, 3)).to(device)
b = torch.normal(0, 1, (2, 3)).to(device) c = a * b在上述位置插入堆栈打印代码并重新编译、安装后运行用例，可以看到以下输出（截取部分）：
usr/local/lib/python3.8/dist-packages/_XLAC.cpython-38-x86_64-linux-gnu.so(_ZN9torch_xla15TensorToXlaDataERKN2at6TensorERKNS_6DeviceEb+0x64d) [0x7f086098b9ed]
/usr/local/lib/python3.8/dist-packages/_XLAC.cpython-38-x86_64-linux-gnu.so(_ZNK9torch_xla9XLATensor19GetIrValueForTensorERKN2at6TensorERKNS_6DeviceE+0xa5) [0x7f0860853955]
/usr/local/lib/python3.8/dist-packages/_XLAC.cpython-38-x86_64-linux-gnu.so(_ZNK9torch_xla9XLATensor10GetIrValueEv+0x19b) [0x7f0860853d5b]
/usr/local/lib/python3.8/dist-packages/_XLAC.cpython-38-x86_64-linux-gnu.so(_ZN9torch_xla9XLATensor3mulERKS0_S2_N3c108optionalINS3_10ScalarTypeEEE+0x3f) [0x7f086087631f]
/usr/local/lib/python3.8/dist-packages/_XLAC.cpython-38-x86_64-linux-gnu.so(_ZN9torch_xla18XLANativeFunctions3mulERKN2at6TensorES4_+0xc4) [0x7f08606d4da4]
/usr/local/lib/python3.8/dist-packages/_XLAC.cpython-38-x86_64-linux-gnu.so(+0x19d158) [0x7f08605f7158]
/usr/local/lib/python3.8/dist-packages/torch/lib/libtorch_cpu.so(_ZN2at4_ops10mul_Tensor10redispatchEN3c1014DispatchKeySetERKNS_6TensorES6_+0xc5) [0x7f0945c9d055]
/usr/local/lib/python3.8/dist-packages/torch/lib/libtorch_cpu.so(+0x2b8986c) [0x7f094705986c]
/usr/local/lib/python3.8/dist-packages/torch/lib/libtorch_cpu.so(+0x2b8a37b) [0x7f094705a37b]
/usr/local/lib/python3.8/dist-packages/torch/lib/libtorch_cpu.so(_ZN2at4_ops10mul_Tensor4callERKNS_6TensorES4_+0x157) [0x7f0945cee717]
/usr/local/lib/python3.8/dist-packages/torch/lib/libtorch_python.so(+0x3ee91f) [0x7f094e4b391f]
/usr/local/lib/python3.8/dist-packages/torch/lib/libtorch_python.so(+0x3eeafb) [0x7f094e4b3afb]
python() [0x5042f9]明显可以看到是从python的堆栈调用过来的，分析一下可以得知_ZN2at4_ops10mul_Tensor10redispatchEN3c1014DispatchKeySetERKNS_6TensorES6_+0xc5对应的定义是at::_ops::mul_Tensor::redispatch(c10::DispatchKeySet, at::Tensor const&, at::Tensor const&)+0xc5
虽然这里意义仍有些不明，但我们已经可以做出推测了：redistpatch函数是根据DispatchKeySet来决定将操作dispatch到某个backend上，xla的device信息就被包含在其中。而后面两个输入的const at::Tensor&就是乘法操作的两个输入。
根据上面的关键字redispatch来寻找，我们可以找到这样一个文件gen.py，其中的codegen函数很多，但最显眼的是下面的OperatorGen：
@dataclass(frozen=True)
class ComputeOperators: target: Union[ Literal[Target.DECLARATION], Literal[Target.DEFINITION] ] @method_with_native_function def __call__(self, f: NativeFunction) -> str: sig = DispatcherSignature.from_schema(f.func) name = f.func.name.unambiguous_name() call_method_name = 'call' redispatch_method_name = 'redispatch' if self.target is Target.DECLARATION: return f"""
struct TORCH_API {name} {{ using schema = {sig.type()}; using ptr_schema = schema*; // See Note [static constexpr char* members for windows NVCC] STATIC_CONSTEXPR_STR_INL_EXCEPT_WIN_CUDA(name, "aten::{f.func.name.name}") STATIC_CONSTEXPR_STR_INL_EXCEPT_WIN_CUDA(overload_name, "{f.func.name.overload_name}") STATIC_CONSTEXPR_STR_INL_EXCEPT_WIN_CUDA(schema_str, {cpp_string(str(f.func))}) static {sig.defn(name=call_method_name, is_redispatching_fn=False)}; static {sig.defn(name=redispatch_method_name, is_redispatching_fn=True)};
}};""" elif self.target is Target.DEFINITION: defns = f"""
STATIC_CONST_STR_OUT_OF_LINE_FOR_WIN_CUDA({name}, name, "aten::{f.func.name.name}")
STATIC_CONST_STR_OUT_OF_LINE_FOR_WIN_CUDA({name}, overload_name, "{f.func.name.overload_name}")
STATIC_CONST_STR_OUT_OF_LINE_FOR_WIN_CUDA({name}, schema_str, {cpp_string(str(f.func))}) // aten::{f.func}
static C10_NOINLINE c10::TypedOperatorHandle<{name}::schema> create_{name}_typed_handle() {{ return c10::Dispatcher::singleton() .findSchemaOrThrow({name}::name, {name}::overload_name) .typed<{name}::schema>();
}}
""" for is_redispatching_fn in [False, True]: if is_redispatching_fn: dispatcher_exprs_str = ', '.join(['dispatchKeySet'] + [a.name for a in sig.arguments()]) dispatcher_call = 'redispatch' method_name = f'{name}::{redispatch_method_name}' else: dispatcher_exprs_str = ', '.join([a.name for a in sig.arguments()]) dispatcher_call = 'call' method_name = f'{name}::{call_method_name}' defns += f"""
// aten::{f.func}
{sig.defn(name=method_name, is_redispatching_fn=is_redispatching_fn)} {{ static auto op = create_{name}_typed_handle(); return op.{dispatcher_call}({dispatcher_exprs_str});
}}
""" return defns else: assert_never(self.target)对于每个算子，PyTorch会（在编译前）在这里生成许多类，这些类会有静态成员call或者redispatch，其中redispatch负责分发具体的实现。这里的codegen比较繁琐，这里就不再细讲。
注册PyTorch库实现
即便我们找到了上面redispatch和codegen的线索，看起来仍然不足以解释PyTorch到torch-xla的桥接，因为PyTorch和torch-xla两个库之间的调用，必须要有符号的映射才可以，而不是一些函数形式上的相同。PyTorch中是有Dispatcher机制的，这个机制很常见于很多框架，比如oneflow也是有一套类似的Dispatcher机制。这套机制最大的好处就是在尽可能减少侵入式修改的前提下保证了较高的可扩展性。简而言之，我们的op有一种定义，但可以有多种实现方式，并且这个实现的代码可以不在框架内部，这样就使得框架在保持通用性的同时，易于在特定环境下做针对性的扩展。这套机制本质上就是建立了一个字典，将op映射到函数指针，那么每次调用一个op的时候，我们可以根据一些标识（比如tensor.device）来判断应该调用哪一种实现。
PyTorch中提供了一个宏用来将实现注册，从而让dispatcher可以调用：
#define _TORCH_LIBRARY_IMPL(ns, k, m, uid) \ static void C10_CONCATENATE( \ TORCH_LIBRARY_IMPL_init_##ns##_##k##_, uid)(torch::Library&); \ static const torch::detail::TorchLibraryInit C10_CONCATENATE( \ TORCH_LIBRARY_IMPL_static_init_##ns##_##k##_, uid)( \ torch::Library::IMPL, \ c10::guts::if_constexpr<c10::impl::dispatch_key_allowlist_check( \ c10::DispatchKey::k)>( \ []() { \ return &C10_CONCATENATE( \ TORCH_LIBRARY_IMPL_init_##ns##_##k##_, uid); \ }, \ []() { return [](torch::Library&) -> void {}; }), \ #ns, \ c10::make_optional(c10::DispatchKey::k), \ __FILE__, \ __LINE__); \ void C10_CONCATENATE( \ TORCH_LIBRARY_IMPL_init_##ns##_##k##_, uid)(torch::Library & m)这个宏如果完全展开会是下面这样：
static void TORCH_LIBRARY_IMPL_init_aten_CPU_0(torch::Library&);
static const torch::detail::TorchLibraryInit TORCH_LIBRARY_IMPL_static_init_aten_CPU_0( torch::Library::IMPL, (c10::impl::dispatch_key_allowlist_check(c10::DispatchKey::CPU) ? &TORCH_LIBRARY_IMPL_init_aten_CPU_0 : [](torch::Library&) -> void {}), "aten", c10::make_optional(c10::DispatchKey::CPU), __FILE__, __LINE__);
void TORCH_LIBRARY_IMPL_init_aten_CPU_0(torch::Library & m)这里比较需要注意的是第二行的TORCH_LIBRARY_IMPL_static_init_aten_CPU_0并不是一个函数，而是一个静态变量，它的作用就是在torch_xla库初始化的时候，将xla定义的op注册到PyTorch中。
关于这部分更详细的介绍可以参考https://zhuanlan.zhihu.com/p/648578629
从PyTorch调用到torch_xla
xla调用上面所说的宏进行注册的位置在RegisterXLA.cpp这个文件中（codegen的结果），如下：
ORCH_LIBRARY_IMPL(aten, XLA, m) { m.impl("abs", TORCH_FN(wrapper__abs)); ...
}其中，wrapper__abs的定义如下：
at::Tensor wrapper__abs(const at::Tensor & self) { return torch_xla::XLANativeFunctions::abs(self);
}显然，这个定义和PyTorch框架内部的算子是完全一致的，只是修改了实现。而XLANativeFunctions::abs的实现可以在aten_xla_type.cpp中找到，如下所示：
at::Tensor XLANativeFunctions::abs(const at::Tensor& self) { XLA_FN_COUNTER("xla::"); return bridge::AtenFromXlaTensor(XLATensor::abs(bridge::GetXlaTensor(self)));
}到这里已经比较明朗了，注册之后，PyTorch上对于op的调用最终会进入torch_xla的native function中调用对应的op实现，而这些实现的根本都是对XLATensor进行操作，在最终操作执行完成之后，会将作为结果的XLATensor重新转换为torch Tensor，但要注意，这里的结果不一定被实际计算了，也可能只是记录了一下IR，将节点加入图中，这取决于具体的实现。
总结
其实torch-xla官方的文档里是有关于代码生成和算子注册这个过程的描述的，只不过一开始我没找到这个文档，走了一点弯路，但是自己探索也会觉得更明了这个过程。官方文档中的描述如下（节选）：
All file mentioned below lives under the xla/torch_xla/csrc folder, with the exception of codegen/xla_native_functions.yaml
1.xla_native_functions.yaml contains the list of all operators that are lowered. Each operator name must directly match a pytorch operator listed in native_functions.yaml. This file serves as the interface to adding new xla operators, and is an input to PyTorch's codegen machinery. It generates the below 3 files: XLANativeFunctions.h, RegisterXLA.cpp, and RegisterAutogradXLA.cpp
2.XLANativeFunctions.h and aten_xla_type.cpp are entry points of PyTorch to the pytorch_xla world, and contain the manually written lowerings to XLA for each operator. XLANativeFunctions.h is auto-generated through a combination of xla_native_functions.yaml and the PyTorch core native_functions.yaml file, and contains declarations for kernels that need to be defined in aten_xla_type.cpp. The kernels written here need to construct 'XLATensor' using the input at::Tensor and other parameters. The resulting XLATensor needs to be converted back to the at::Tensor before returning to the PyTorch world
3.RegisterXLA.cpp and RegisterAutogradXLA.cpp are auto-generated files that register all lowerings to the PyTorch Dispatcher. They also include auto-generated wrapper implementations of out= and inplace operators.
大概意思就是实际上torch-xla就是根据xla_native_functions.yaml这个文件来生成算子的定义，然后再生成对应的RegisterXLA.cpp中的注册代码，这也跟PyTorch的codegen方式一致。
综合这一整个过程可以看出，PyTorch是保持了高度的可扩展性的，不需要多少侵入式的修改就可以将所有的算子全部替换成自己的，这样的方式也可以让开发者不用去关注dispatcher及其上层的实现，专注于算子本身的逻辑。
二、PyTorch语义分割实现洪水识别
全世界有数百万人因洪水而流离失所。确实，通过使用深度学习 + 计算机视觉，我们无法总是预测下一次洪水何时会来袭。但我们可以在洪水灾区的图像上训练语义分割算法。这样的模型可以帮助分析和决策未来的情况。为了尽我们的微薄之力，我们将在本文中使用 PyTorch训练一个用于洪水识别的语义分割模型。
洪水分割数据集
洪水分割看起来可能与卫星图像中的水体分割类似。我们在尝试在自定义数据集上训练 PyTorch DeepLabV3 模型时涵盖了水体分割项目。
但在尝试划分洪水区域时，情况会变得更加复杂。让我们先探索数据集，然后再进一步讨论。
我们将使用 Kaggle 的洪水区域分割的修改版本。这是一个非常好且干净的数据集。但它存在一些图像扩展问题，并且其中一个图像已损坏。
https://www.kaggle.com/datasets/faizalkarim/flood-area-segmentation我准备了这个数据集的修改版本。这是Kaggle 上的洪水分割数据集训练/验证分割。我们将使用此版本的数据集。以下是更改：
- 数据集现在包含训练和验证部分。
- 有 257 张训练图像和蒙版以及 32 张验证图像和蒙版。
- 现在数据集中没有损坏的图像。
下载并提取数据集后，您将发现以下结构。
flood-area-segmentation
├── train_images
├── train_masks
├── valid_images
└── valid_masks所有地面实况图像均为 RGB 格式。所有蒙版均为灰度图像，洪水区域为白色像素（255），其他区域为黑色像素（0）。因此，这是一个二元语义分割数据集。以下是一些图像。
洪水识别的语义分割为何如此困难？
让我们来谈谈为什么根据这个数据集进行洪水分割很困难。
首先，这个数据集仅包含 257 张训练图像和掩码。当然，这不足以训练最先进的语义分割模型。我们需要足够的增强技术才能减少损失。
其次，洪水区域之间会有很多小的背景像素，包括房屋、农田，甚至是部分被淹没的车辆。
以下列举几个相同的例子。
模型必须正确学习这些背景类别，而不是对其进行分割。对于较小的模型来说，这可能特别具有挑战性。
考虑到上述问题，我们将尽力训练出最好的模型。
PyTorch安装与项目目录
该项目的代码使用TORCH 1.12.0 和 TORCHVISION 0.13.0。您也可以使用较新的版本。
此外，Albumentations是图像和蒙版增强的要求之一。
请务必安装以上两个库。
洪水识别语义分割的目录结构：
.
├── input
│ ├── flood-area-segmentation
│ │ ├── train_images
│ │ ├── train_masks
│ │ ├── valid_images
│ │ └── valid_masks
│ └── inference_data
│ └── video_1.mp4
├── outputs
│ ├── inference_results
│ │ ├── 1005.jpg
│ │ ...
│ │ └── 38.jpg
│ ├── inference_results_video
│ │ └── video_1.mp4
│ ├── valid_preds
│ │ ├── e0_b7.jpg
│ │ ...
│ │ └── e9_b7.jpg
│ ├── accuracy.png
│ ├── best_model.pth
│ ├── loss.png
│ └── model.pth
└── src├── config.py├── datasets.py├── engine.py├── inference_image.py├── inference_video.py├── metrics.py├── model.py├── train.py└── utils.pyinput目录包含洪水区域分割数据集和推理数据。
outputs目录中，我们拥有所有与训练和推理相关的输出。这包括模型权重文件、推理后的图像结果以及验证循环的分割输出。
src目录包含所有 Python 源代码文件。由于有 9 个 Python 文件，我们无法详细介绍每个文件。但我们将在下一节中探讨代码的重要部分。
使用DeepLabV3 ResNet50 进行洪水识别的语义分割
从现在开始，我们将讨论项目的技术方面。在开始训练之前，我们将讨论以下主题：
- 配置文件。
- 用于洪水识别语义分割的深度学习模型。
- 数据集准备策略。
- 训练集的数据增强。
- 训练超参数和参数。这些包括优化器、学习率以及根据 VRAM 可用性选择批处理大小。我们还将讨论训练图像大小的影响。
配置文件
我们将在其余 Python 文件中使用一些配置。我们不必反复重写这些配置，只需config.py归档并存储这些。
ALL_CLASSES = ['background', 'flood']
LABEL_COLORS_LIST = [(0, 0, 0), # Background.(255, 255, 255), # Flood.
]
VIS_LABEL_MAP = [(0, 0, 0), # Background.(255, 0, 0), # Flood.
]ALL_CLASSED列表包含数据集中的类的名称。
LABEL_COLORS_LIST列表包含数据集中的 RGB 颜色元组。在数据集中，背景类别为黑色，淹没区域为白色。
最后，VIS_LABEL_MAP包含我们将用于可视化的 RGB 颜色元组。为了清晰起见，我们将在推理过程中使用红色来划分被淹没的区域。
DeepLabV3 ResNet50 模型
我们将在此项目中使用优质、可靠、经典的DeepLabV3 和 ResNet50主干进行语义分割。
准备用于自定义模型训练的模型非常简单。事实上，我们只需使用以下几行代码即可准备用于自定义语义分割训练的模型。
import torch.nn as nn
from torchvision.models.segmentation import deeplabv3_resnet50def prepare_model(num_classes=2):model = deeplabv3_resnet50(weights='DEFAULT')model.classifier[4] = nn.Conv2d(256, num_classes, 1)model.aux_classifier[4] = nn.Conv2d(256, num_classes, 1)return model为了获得更好的性能，我们将使用预训练的权重。另外，我们需要修改分类器和辅助分类器根据数据集中的类别数量，将 head 分为两类：背景类和洪水泛滥区类。
选择带有 ResNet 的 DeepLabV3 的主要原因是速度和性能。在 PyTorch 中，我们还有另外两个 DeepLabV3 选项：
- 采用 ResNet101 主干。虽然性能良好，但需要更多内存进行训练，推理速度也较慢。
- 还有以 MobileNet 为骨干的 DeepLabV3。速度很快，但性能不太好。
为此，我们选择 DeepLabV3 ResNet50 来兼顾两全其美。
数据集准备
虽然我们不会介绍数据集准备的整个代码，但这里有几件重要的事情需要注意。
其中之一就是数据增强策略。
请记住，我们只有 257 张图像和蒙版用于训练。这不足以训练出一个好的模型。但数据增强可以帮助我们做到这一点。我们将使用 Albumentations 库轻松地将增强功能应用于图像和蒙版。
我们将对洪水分割数据集的训练图像和蒙版应用以下增强功能。
- 水平翻转
- 隨機亮度對比
- 随机太阳光晕
- 随机雾
- 图像压缩
除了水平翻转（以 0.5 的概率应用）之外，我们将以 0.2 的概率应用所有其他增强。
这是随机应用这些增强功能后图像的样子。
如您所见，仅对蒙版应用了空间增强（翻转）。影响颜色、色调或饱和度的其他像素级操作不应用于蒙版。使用 Albumentations 很容易实现这一点。请查看dataset.py文件以获取更多详细信息。
使用脚本、指标以及训练和验证代码
utils.py文件包含许多辅助函数和实用程序类。其中包括用于保存最佳模型的类、用于保存准确率和损失图的函数，以及用于在图像上叠加分割图的函数。
metrics.py包含pix_acc函数。这将返回标记像素的总数和正确分类的像素。我们将使用这些来计算engine.py文件。
因此，engine.py包含我们将在每个时期调用的训练和验证函数。
在洪水分割数据集上训练 DeepLabV3
最后，我们可以在洪水分割数据集上训练 DeepLabV3 ResNet50 模型了。
注意：所有训练和推理实验均在具有 10GB RTX 3080 GPU、第 10 代 i7 CPU 和 32 GB RAM 的机器上进行。
train.py是我们将用来启动训练的驱动脚本。
在开始训练之前，让我们先来看看训练脚本支持的参数解析器。
--epoches：我们想要训练模型的时期数。
--lr：优化器的学习率，默认为0.0001。
--batch：数据加载器的批次大小。
--imgsz：我们想要用来训练模型的图像大小。语义分割模型通常需要更高分辨率的图像才能取得良好的效果。我们将在 512×512 图像上训练模型，这也是默认值。
--scheduler：这是一个布尔标志，表示我们是否要使用学习率调度程序。如果我们传递此标志，则学习率将在 25 个时期后降低 10 倍。
如果在训练时遇到 OOM（内存不足）错误，请考虑减少图像大小或批量大小。
我们可以通过在终端中执行以下命令来开始训练源码作为当前工作目录。
python train.py --epochs 50 --batch 4 --lr 0.0001 --scheduler我们对模型进行 50 个epoch的训练，batch size为 4，起始学习率为 0.0001，并应用学习率调度程序。
以下是终端的截断输出。
python train.py --epochs 50 --batch 4 --lr 0.0001 --scheduler
Namespace(epochs=50, lr=0.0001, batch=4, scheduler=True)
41,994,308 total parameters.
41,994,308 training parameters.
Adjusting learning rate of group 0 to 1.0000e-04.
EPOCH: 1
Training
Loss: 0.4229 | PixAcc: 83.23: 100%|████████████████████| 64/64 [00:50<00:00, 1.26it/s]
Validating
Loss: 0.3060 | PixAcc: 89.33: 100%|████████████████████| 8/8 [00:02<00:00, 3.57it/s]
Best validation loss: 0.38311174884438515
Saving best model for epoch: 1
Train Epoch Loss: 0.5018, Train Epoch PixAcc: 64.9515
Valid Epoch Loss: 0.3831, Valid Epoch PixAcc: 86.8635
Adjusting learning rate of group 0 to 1.0000e-04.
--------------------------------------------------
EPOCH: 2
Training
Loss: 0.3776 | PixAcc: 81.60: 100%|████████████████████| 64/64 [00:49<00:00, 1.30it/s]
Validating
Loss: 0.4301 | PixAcc: 91.83: 100%|████████████████████| 8/8 [00:02<00:00, 3.84it/s]
Train Epoch Loss: 0.3887, Train Epoch PixAcc: 72.3201
Valid Epoch Loss: 0.4537, Valid Epoch PixAcc: 94.0050
Adjusting learning rate of group 0 to 1.0000e-04.
--------------------------------------------------
EPOCH: 3
Training
Loss: 0.4062 | PixAcc: 68.83: 100%|████████████████████| 64/64 [00:49<00:00, 1.29it/s]
Validating
Loss: 0.2974 | PixAcc: 93.10: 100%|████████████████████| 8/8 [00:02<00:00, 3.86it/s]
Train Epoch Loss: 0.3698, Train Epoch PixAcc: 75.0250
Valid Epoch Loss: 0.4060, Valid Epoch PixAcc: 94.7763
Adjusting learning rate of group 0 to 1.0000e-04.
--------------------------------------------------
.
.
.
EPOCH: 50
Training
Loss: 0.1831 | PixAcc: 80.59: 100%|████████████████████| 128/128 [00:39<00:00, 3.26it/s]
Validating
Loss: 0.1339 | PixAcc: 98.66: 100%|████████████████████| 16/16 [00:02<00:00, 7.16it/s]
Train Epoch Loss: 0.1459, Train Epoch PixAcc: 91.4595
Valid Epoch Loss: 0.2437, Valid Epoch PixAcc: 88.4480
Adjusting learning rate of group 0 to 1.0000e-06.
--------------------------------------------------
TRAINING COMPLETE分析 DeepLabV3 洪水分割结果
有趣的是，我们在 epoch 33 上获得了最小的验证损失 0.222，而最高的像素准确度出现在 epoch 3。这种情况有点奇怪，因为最佳模型是基于最小损失保存的。让我们来看看准确度和损失图。
我们可以看到在初始阶段和整个训练过程中出现了很多波动，直到学习率下降。这意味着我们需要学习率调度程序。也许我们可以用更低的初始学习率进行训练，以缓解这种不稳定的训练。
无论如何，我们现在有一个训练好的模型。让我们对图像和视频进行一些推理实验。
图像洪水分割推理
对于图像推理，我们将从数据集中选择相同的验证图像。
为了对图像进行推理，我们将使用inference_image.py脚本。
python inference_image.py --input ../input/flood-area-segmentation/valid_images/我们使用--input标志提供验证图像目录的路径。
您可以在里面找到outputs/inference_results目录。
以下是一些输出。第一列显示模型的预测，右列显示真实值掩码。
虽然结果可能并不完美，但效果相当不错。该模型仍然难以区分细长的陆地线和被淹没的区域。但我们可以通过更多的训练数据和更好的训练技术来纠正这个问题。
视频洪水分割推理
python inference_video.py --input ../input/inference_data/ video.mp4在 RTX 3080 GPU 上，我们平均获得 21 FPS。这可能并不像你想象的那么好。我们在将帧输入模型之前不会调整它们的大小。将它们调整为 512×512 分辨率将稍微提高 FPS。
以下是我们得到的输出。
总体来说，效果还不错。如果将其部署到无人机上，即可实时观察洪水情况了。
源码下载：
链接: https://pan.baidu.com/s/1Dg9q3ofvUKlm5AAYnoNnJA 提取码: z7u3数据集下载：
https://www.kaggle.com/datasets/faizalkarim/flood-area-segmentation三、PyTorch~Async Checkpoint Save
本文介绍了PyTorch 2.4的新特性——异步Checkpoint保存功能，该功能通过将模型状态的保存过程移到CPU线程来减少训练中断，大幅提升了大型模型训练时Checkpoint保存的速度。
PyTorch Async Checkpoint Save
PyTorch博客资料：https://pytorch.org/blog/reducing-checkpointing-times/
PyTorch实现和使用Demo：https://github.com/pytorch/pytorch/blob/main/torch/distributed/checkpoint/state_dict_saver.py
功能介绍
在PyTorch 2.4之后，我们可以尝试使用PyTorch开发的异步Checkpoint保存功能，这个功能是和IBM联合开发的，在7B的大模型训练中，Checkpoint保存的时间从平均 148.8 秒缩短至 6.3 秒，快了 23.62 倍。这可以转化为以下两种好处：
- 在继续鲁棒的保存Checkpoint的同时，在每给定的 24 小时内实现更多净训练进度；
- 可以更频繁地进行Checkpoint保存以缩短训练恢复窗口或者时间。
从结果图来看，无论是单机的FSDP还是多机的HSDP，Async Checkpoint Save都展现出了很大的速度优势，对于参数量更大的模型预计收益会更大。目前在TorchTian（https://github.com/pytorch/torchtitan）中已经集成了这个新的功能，相信其他的主流训练框架都会很快跟进此feature。
背景
模型Checkpoint是大模型训练的重要组成部分，但Checkpoint是一个昂贵的过程，因为每个Checkpoint过程都需要阻止训练进度以保存最新的模型权重。但是，不进行Checkpoint或降低Checkpoint频率会导致训练进度损失很大。例如，死锁、straggler（落后者）和 GPU 错误等故障导致需要重新启动训练过程。为了从故障中重启，所有（训练）工作者必须停止其训练过程并从上次保存的Checkpoint重新启动。
因此，对故障的鲁棒性与训练进度之间很难做到权衡，但现在有了异步Checkpoint，PyTorch 分布式训练能够显著缓解这种张力，并以最小的影响整体训练时间的方式实现频繁Checkpoint。
大约一年前(https://pytorch.org/blog/performant-distributed-checkpointing/)，我们展示了分布式Checkpoint如何大幅加速Checkpoint时间，从最初的 torch.save() 功能开始。正如 IBM 研究团队指出的那样，torch.save 可能需要长达 30 分钟才能检查一个 11B 模型（PyTorch 1.13）。
随着分布式Checkpoint的进步，对于高达 30B 的模型大小，Checkpoint可以在 4 分钟内完成。使用异步Checkpoint，Checkpoint导致的训练时间损失现在降至 30 秒以下，通常仅需 6 秒。
需要明确的是，异步Checkpoint不会压缩实际的序列化Checkpoint时间，如之前的更新所展示的那样。相反，它将最终的Checkpoint过程移出关键路径（到 CPU 线程），以允许 GPU 训练在单独的线程下完成Checkpoint的同时继续进行。
如上图所示，异步Checkpoint比一年前的改进进一步提高了 10 倍到 23 倍。
Async Checkpoint Save如何工作
异步Checkpoint将Checkpoint过程模块化分为两个部分，而不是一个单一的整体过程。第一阶段将每个 GPU/rank 的数据从 GPU 复制到 CPU。这是用户可见的停机时间，对于 7B-13B 的模型大小可能需要 6 到 14 秒。第二阶段异步地将数据从 CPU 内存复制到磁盘以持久保存Checkpoint。
一旦数据在第一阶段复制到 CPU，GPU 就可以立即恢复训练。因此，使用异步Checkpoint，Checkpoint的停机时间仅仅是将最新的模型状态复制到 CPU 所需的时间。在训练恢复的同时，非阻塞 CPU 线程使用内存中新到达的数据完成完整的Checkpoint/序列化过程到磁盘（即持久保存）。
注意，PyTorch 的分布式Checkpoint依赖于集合通信调用来获取必要的每个等级元数据以优化保存，以及最终的同步，该同步将Checkpoint标记为已完成并使操作成为原子操作。如果Checkpoint线程使用与训练相同的进程组，这可能会干扰分布式训练（因为分布式训练也依赖于类似的调用来跨多个 GPU 同步训练）。具体来说，调用之间的竞争条件可能会导致训练和异步Checkpoint保存线程同时等待集合通信调用，从而导致真正的集合通信卡死。我们通过为异步Checkpoint初始化一个单独的进程组来避免这种情况。这将Checkpoint集合通信分离到其自己的逻辑进程组中，从而确保它不会干扰主训练线程中的集合通信调用。
如何使用PyTorch Async Checkpoint Save
这里是最小的使用PyTorch Async Checkpoint Save的demo：
需要注意的是第12行，为异步的Checkpoint集合通信操作建立了一个新的group，然后调用dcp.save的时候我们需要传入这个group。
https://github.com/pytorch/torchtitan 里面也已经使用上了这个功能，可以用于预训练自己的 Llama2 或 Lllama3 模型。在配置文件里面就可以选择使用Async Checkpoint Save。如下图所示：
代码流程粗略浏览
代码实现在 https://github.com/pytorch/pytorch/blob/main/torch/distributed/checkpoint/state_dict_saver.py 这个文件中。核心部分为以下2个函数，这里简单展示下流程：
# 创建 state_dict 的浅拷贝，对于每个 Stateful 对象调用其 state_dict() 方法。
def _stateful_to_state_dict(state_dict: STATE_DICT_TYPE) -> STATE_DICT_TYPE: """Creates a shallow copy of `state_dict` where `state_dict` is called for each Stateful object.""" stateful_state_dict = {} for key, elem in state_dict.items(): stateful_state_dict[key] = ( elem.state_dict() if isinstance(elem, Stateful) else elem ) return stateful_state_dict @_dcp_method_logger(log_exceptinotallow=True)
def async_save( state_dict: STATE_DICT_TYPE, *, checkpoint_id: Union[str, os.PathLike, None] = None, storage_writer: Optional[StorageWriter] = None, planner: Optional[SavePlanner] = None, process_group: Optional[dist.ProcessGroup] = None,
) -> Future: torch._C._log_api_usage_once("torch.distributed.checkpoint.async_save") # 检查分布式环境设置 if dist.is_available() and dist.is_initialized(): pg = process_group or _get_default_group() assert ( torch.device("cpu") in pg._device_types # type: ignore[attr-defined] ), "A CPU backend must be enabled for async save; try initializing process group with 'cpu:gloo,cuda:nccl'" # 设置存储写入器 storage_writer = cast( StorageWriter, _storage_setup(storage_writer, checkpoint_id, reader=False) ) # 处理状态字典（调用 _stateful_to_state_dict） state_dict = _stateful_to_state_dict(state_dict) # 如果存储写入器支持异步暂存，则使用它；否则将状态字典卸载到 CPU if isinstance(storage_writer, AsyncStager): staged_state_dict = storage_writer.stage(state_dict) else: # provides bwc for storage_writers not implementing AsyncStager staged_state_dict = _offload_state_dict_to_cpu(state_dict, type_check=False) # 创建线程池执行器，提交保存任务。这里是一个线程 executor = ThreadPoolExecutor(max_workers=1) f: Future = executor.submit( save, staged_state_dict, checkpoint_id=checkpoint_id, storage_writer=storage_writer, planner=planner, process_group=process_group, ) # 设置任务完成后的回调函数（关闭执行器） f.add_done_callback(lambda f: executor.shutdown(wait=False)) # 如果需要，同步暂存操作 if ( isinstance(storage_writer, AsyncStager) and storage_writer.should_synchronize_after_execute ): storage_writer.synchronize_staging() # 返回 Future 对象 return f将来的改进
PyTorch Blog中提到Checkpoint在过去的一年中取得了巨大进步。从近半个小时的Checkpoint变为使用分布式Checkpoint不到 5 分钟，现在又变为使用异步Checkpoint不到 30 秒。最后一个前沿是零开销Checkpoint，即使是小于 30 秒的时间也可以通过在反向传递期间流式传输更新的权重来消除，这样Checkpoint数据在异步Checkpoint开始时就已经在 CPU 上了。这将有效地将大型模型训练转移到Checkpoint没有中断或停机时间的程度，从而既提高了鲁棒性（因为可以更频繁地进行Checkpoint），又因为没有Checkpoint的停机时间而加快了训练进度。
四、Pytorch量化新方法TorchAO
Pytorch的量化方法切换到了torchao，本篇基于官方教程简单介绍下torchao的量化使用教程。
使用 TorchAO 实现 GPU 量化
本篇对segment anything 模型进行量化和优化。参考了 segment-anything-fast 仓库时所采取的步骤。
本指南演示了如何应用这些技术来加速模型，尤其是那些使用 Transformer 的模型。为此，我们将重点关注广泛适用的技术，例如使用 torch.compile 进行性能优化和量化，并衡量其影响。
环境
实验环境：
- CUDA 12.1
- A100-PG509-200，功率限制为 330.00 W
不同硬件可能结果不同。
conda create -n myenv python=3.10
pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu121
pip install git+https://github.com/facebookresearch/segment-anything.git
pip install git+https://github.com/pytorch-labs/ao.gitSegment Anything Model checkpoint：
访问 segment-anything checkpoint，并下载 `vit_h` checkpoint。或者可以使用 `wget` （例如，`wget https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth --directory-prefix=<path>`）。通过编辑以下代码传入该目录路径：sam_checkpoint_base_path = <path>import torch
from torchao.quantization import change_linear_weights_to_int8_dqtensors
from segment_anything import sam_model_registry
from torch.utils.benchmark import Timer sam_checkpoint_base_path = "data"
model_type = 'vit_h'
model_name = 'sam_vit_h_4b8939.pth'
checkpoint_path = f"{sam_checkpoint_base_path}/{model_name}"
batchsize = 16
only_one_block = True @torch.no_grad()
def benchmark(f, *args, **kwargs): for _ in range(3): f(*args, **kwargs) torch.cuda.synchronize() torch.cuda.reset_peak_memory_stats() t0 = Timer( stmt="f(*args, **kwargs)", globals={"args": args, "kwargs": kwargs, "f": f} ) res = t0.adaptive_autorange(.03, min_run_time=.2, max_run_time=20) return {'time': res.median * 1e3, 'memory': torch.cuda.max_memory_allocated() / 1e9} def get_sam_model(only_one_block=False, batchsize=1): sam = sam_model_registry[model_type](checkpoint=checkpoint_path).cuda() model = sam.image_encoder.eval（) image = torch.randn(batchsize, 3, 1024, 1024, device='cuda') # 使用模型的单个 block if only_one_block: model = model.blocks[0] image = torch.randn(batchsize, 64, 64, 1280, device='cuda') return model, image在本教程中，我们将重点量化 image_encoder，因为它的输入是固定尺寸的，而 prompt encoder 和 mask decoder 的尺寸是dynamic，量化这些模块更为复杂。
我们首先从单个 block 开始，以简化分析。
基准测试
首先，我们来测量模型的基础运行时间：
try: model, image = get_sam_model(only_one_block, batchsize) fp32_res = benchmark(model, image) print(f"模型的基础 fp32 运行时间为 {fp32_res['time']:0.2f}ms，峰值内存为 {fp32_res['memory']:0.2f}GB")
except Exception as e: print("无法运行 fp32 模型：", e)模型的基础 fp32 运行时间为 198.00ms，峰值内存为 8.54GB。
使用 bfloat16 提升性能
通过将模型转换为 bfloat16 格式，直接就有性能提升。我们选择 bfloat16 而不是 fp16 的原因是它的动态范围与 fp32 相当。bfloat16 和 fp32 具有相同的 8 位指数，而 fp16 只有 4 位。较大的动态范围有助于防止溢出错误及其他可能因量化而出现的问题。
model, image = get_sam_model(only_one_block, batchsize)
model = model.to(torch.bfloat16)
image = image.to(torch.bfloat16)
bf16_res = benchmark(model, image)
print(f"bf16 block 的运行时间为 {bf16_res['time']:0.2f}ms，峰值内存为 {bf16_res['memory']: 0.2f}GB")bf16 block 的运行时间为 70.45ms，峰值内存为 5.38GB。
通过此简单的更改，运行时间提高了约 7 倍（从 186.16ms 到 25.43ms）。
使用 torch.compile 进行编译优化
接下来，我们使用 torch.compile 对模型进行编译，看看性能有多大提升：
model_c = torch.compile(model, mode='max-autotune')
comp_res = benchmark(model_c, image)
print(f"bf16 编译后 block 的运行时间为 {comp_res['time']:0.2f}ms，峰值内存为 {comp_res['memory']: 0.2f}GB")torch.compile 提供了大约 27% 的性能提升。
量化
接下来，我们将应用量化。对于 GPU，量化主要有三种形式：
- int8 动态量化
- int8 仅权重量化
- int4 仅权重量化
不同的模型或模型中的不同层可能需要不同的量化技术。在此示例中，Segment Anything 模型是计算密集型的，因此我们使用动态量化：
del model_c, model, image
model, image = get_sam_model(only_one_block, batchsize)
model = model.to(torch.bfloat16)
image = image.to(torch.bfloat16)
change_linear_weights_to_int8_dqtensors(model)
model_c = torch.compile(model, mode='max-autotune')
quant_res = benchmark(model_c, image)
print(f"bf16 量化后 block 的运行时间为 {quant_res['time']:0.2f}ms，峰值内存为 {quant_res['memory']: 0.2f}GB")通过量化，我们进一步提高了性能，但内存使用显著增加。
内存优化
我们通过融合整数矩阵乘法与后续的重新缩放操作来减少内存使用：
torch._inductor.config.force_fuse_int_mm_with_mul = True通过这种方式，我们再次提升了性能，且大大减少了内存的增长。
进一步优化
最后，我们还可以应用一些通用的优化措施来获得最终的最佳性能：
- 禁用 epilogue fusion
- 应用坐标下降优化
torch._inductor.config.epilogue_fusion = False
torch._inductor.config.coordinate_descent_tuning = True
torch._inductor.config.coordinate_descent_check_all_directions = True总结
通过本教程，我们了解了如何通过量化和优化技术加速 Segment Anything 模型。在批量大小为 16 的情况下，最终模型量化加速大约为 7.7%。
五、PyTorch~depyf：轻松掌握 torch.compile
本文介绍了 depyf 工具库，旨在帮助开发者理解和掌握 PyTorch 的 torch.compile 特性，通过生成易于阅读的调试信息，降低了学习曲线，使优化过程更加透明。同时，展示了如何使用 depyf 逐步调试并分析 torch.compile 优化后的代码。
博客链接：https://pytorch.org/blog/introducing-depyf/
最近了解
torch.compile的时候，发现清华推出了一个可以帮助我们理解torch.compile到底对我们的代码做了什么优化的库depyf，这篇教程是这个库的一个简要介绍，前面对这个教程做了一个翻译。后面一部分，我利用cursor来完整展示了如何完整的阅读depfy生成的torch.compile编译产物的例子，我们可以看到torch.compile优化的每个子图以及产生的fuse kernel，希望对感兴趣的读者有帮助。
介绍 depyf：轻松掌握 torch.compile
很高兴介绍 depyf，这是 PyTorch 生态系统中的一个新项目，旨在帮助用户理解、学习和适应 torch.compile ！
动机
torch.compile 是 PyTorch 2.x 的一个基石，为加速机器学习工作流程提供了一个直接的途径，只需一行代码就可以同时用于训练和推理。仅仅包含 @torch.compile 就可以显著提升你的代码性能。然而，找到 torch.compile 的最佳插入点并不容易，更不用说为了最大效率而调整各种参数的复杂性。
torch.compile 技术栈的复杂性，包括 Dynamo、AOTAutograd、Inductor 等，呈现出一个陡峭的学习曲线。这些对深度学习性能优化至关重要的组件，如果没有坚实的基础知识，可能会令人望而生畏。
注：关于
torch.compile 工作原理的入门示例，请参阅这个逐步说明(https://depyf.readthedocs.io/en/latest/walk_through.html)。
一个常用工具：TORCH_COMPILE_DEBUG
为了揭开 torch.compile 的神秘面纱，常用的方法是利用 TORCH_COMPILE_DEBUG 环境变量。虽然它提供了更多信息，但解读输出仍然是一项艰巨的任务。
例如，当我们有以下代码：
# test.py
import torch
from torch import _dynamo as torchdynamo
from typing import List @torch.compile
def toy_example(a, b): x = a / (torch.abs(a) + 1) if b.sum() < 0: b = b * -1 return x * b def main(): for _ in range(100): toy_example(torch.randn(10), torch.randn(10)) if __name__ == "__main__": main()当我们用 TORCH_COMPILE_DEBUG=1 python test.py 运行它时，我们会得到一个名为 torch_compile_debug/run_2024_02_05_23_02_45_552124-pid_9520 的目录，其中包含这些文件：
.
├── torchdynamo
│ └── debug.log
└── torchinductor ├── aot_model___0_debug.log ├── aot_model___10_debug.log ├── aot_model___11_debug.log ├── model__4_inference_10.1 │ ├── fx_graph_readable.py │ ├── fx_graph_runnable.py │ ├── fx_graph_transformed.py │ ├── ir_post_fusion.txt │ ├── ir_pre_fusion.txt │ └── output_code.py ├── model__5_inference_11.2 │ ├── fx_graph_readable.py │ ├── fx_graph_runnable.py │ ├── fx_graph_transformed.py │ ├── ir_post_fusion.txt │ ├── ir_pre_fusion.txt │ └── output_code.py └── model___9.0 ├── fx_graph_readable.py ├── fx_graph_runnable.py ├── fx_graph_transformed.py ├── ir_post_fusion.txt ├── ir_pre_fusion.txt └── output_code.py生成的文件和日志常常引发的问题比它们解答的还多，让开发者对数据的含义和关系感到困惑。TORCH_COMPILE_DEBUG 的常见疑问包括：
-
model__4_inference_10.1 是什么意思？ - 我只有一个函数，但目录中有三个
model__xxx.py，它们之间有什么对应关系？ -
debug.log 中那些 LOAD_GLOBAL 是什么东西？
更好的工具：DEPYF 来救援
让我们看看 depyf 如何帮助开发者解决上述挑战。要使用 depyf，只需执行 pip install depyf 或按照项目页面 https://github.com/thuml/depyf 安装最新版本，然后用 with depyf.prepare_debug 包围主代码。
# test.py
import torch
from torch import _dynamo as torchdynamo
from typing import List @torch.compile
def toy_example(a, b): x = a / (torch.abs(a) + 1) if b.sum() < 0: b = b * -1 return x * b def main(): for _ in range(100): toy_example(torch.randn(10), torch.randn(10)) if __name__ == "__main__": import depyf with depyf.prepare_debug("depyf_debug_dir"): main()执行 python test.py 后，depyf 将生成一个名为 depyf_debug_dir（prepare_debug 函数的参数）的目录。在该目录下，会有这些文件：
.
├── __compiled_fn_0 AFTER POST GRAD 0.py
├── __compiled_fn_0 Captured Graph 0.py
├── __compiled_fn_0 Forward graph 0.py
├── __compiled_fn_0 kernel 0.py
├── __compiled_fn_3 AFTER POST GRAD 0.py
├── __compiled_fn_3 Captured Graph 0.py
├── __compiled_fn_3 Forward graph 0.py
├── __compiled_fn_3 kernel 0.py
├── __compiled_fn_4 AFTER POST GRAD 0.py
├── __compiled_fn_4 Captured Graph 0.py
├── __compiled_fn_4 Forward graph 0.py
├── __compiled_fn_4 kernel 0.py
├── __transformed_code_0_for_torch_dynamo_resume_in_toy_example_at_8.py
├── __transformed_code_0_for_toy_example.py
├── __transformed_code_1_for_torch_dynamo_resume_in_toy_example_at_8.py
└── full_code_for_toy_example_0.py这里有两个明显的好处：
- 冗长且难以理解的
torchdynamo/debug.log 不见了。它的内容被整理并以人类可读的源代码形式显示在 full_code_for_xxx.py 和 _transformed_code{n}_for_xxx.py 中。值得注意的是，depyf 最艰巨和困难的任务是将 torchdynamo/debug.log 中的字节码反编译成 Python 源代码，从而使开发者免于被 Python 内部结构所困扰。 - 函数名称与计算图之间的对应关系得到了保留。例如，在
__transformed_code_0_for_toy_example.py 中，我们可以看到一个名为 __compiled_fn_0 的函数，我们立即就知道它对应的计算图在 __compiled_fn_0_xxx.py 中，因为它们共享相同的 __compiled_fn_0 前缀名称。
从 full_code_for_xxx.py 开始，并跟随涉及的函数，用户将清楚地了解 torch.compile 对他们的代码做了什么。
再补充一点：逐步调试功能
使用调试器逐行步进代码是理解代码工作原理的好方法。然而，在 TORCH_COMPILE_DEBUG 模式下，这些文件仅供用户参考，无法与用户关心的数据一起执行。
注：这里的"调试"指的是检查和改进程序的过程，而不是纠正有问题的代码。
depyf 的一个突出特点是它能够为 torch.compile 提供逐步调试功能：它生成的所有文件都与 Python 解释器内部的运行时代码对象链接，我们可以在这些文件中设置断点。使用方法很简单，只需添加一个上下文管理器 with depyf.debug()，它就能发挥作用。
# test.py
import torch
from torch import _dynamo as torchdynamo
from typing import List @torch.compile
def toy_example(a, b): x = a / (torch.abs(a) + 1) if b.sum() < 0: b = b * -1 return x * b def main(): for _ in range(100): toy_example(torch.randn(10), torch.randn(10)) if __name__ == "__main__": import depyf with depyf.prepare_debug("depyf_debug_dir"): main() with depyf.debug(): main()需要注意的一点是：调试 torch.compile 的工作流程与标准调试工作流程有所不同。使用 torch.compile 时，许多代码是动态生成的。因此，我们需要：
- 启动程序
- 当程序退出
with depyf.prepare_debug("depyf_debug_dir") 时，代码将在 depyf_debug_dir 中可用。 - 当程序进入
with depyf.debug() 时，它会在内部自动设置一个断点，使程序暂停。 - 导航到
depyf_debug_dir 设置断点。 - 继续运行代码，调试器将会命中这些断点！
这是它看起来的样子的截图。所有代码和张量变量都是实时的，我们可以检查任何变量，并像日常调试工作流程一样逐步执行代码！唯一的区别是我们正在调试 torch.compile 生成的代码，而不是人工编写的代码。
结论
torch.compile 是一个无价的工具，可以轻松加速 PyTorch 代码。对于那些希望深入研究 torch.compile 的人来说，无论是为了充分发挥其潜力还是集成自定义操作，学习曲线可能会非常陡峭。depyf 旨在降低这个障碍，提供一个用户友好的体验来理解、学习和适应 torch.compile。
请亲自探索 depyf 并体验其好处！该项目是开源的，可以在 https://github.com/thuml/depyf 上获取。安装非常简单，只需通过 pip install depyf 即可。我们希望 depyf 能够改善每个人使用 torch.compile 的开发工作流程。
个人补充：用cursor辅助阅读depyf code
使用上面的例子：
# test.py
import torch
from torch import _dynamo as torchdynamo
from typing import List @torch.compile
def toy_example(a, b): x = a / (torch.abs(a) + 1) if b.sum() < 0: b = b * -1 return x * b def main(): for _ in range(100): toy_example(torch.randn(10), torch.randn(10)) if __name__ == "__main__": main()可以顺利得到下面的：
从full_code_for_toy_example_0.py中我们就可以来把握torch.compile到底对我们的代码做了什么优化。这个过程中我们还可以利用LLM加速我们理解代码的过程。例如，在Cursor中，我选中这个文件后输入：可以根据这个文件的内容帮我画一个流程图出来吗？它就会根据这个代码生成一个流程图出来，输出如下：
开始 | v
[输入 a 和 b] | v
[检查 guard 条件] | +-> 如果 guard 条件满足: | | | v | [执行 __transformed_code_0_for_toy_example] | | | v | [计算 x = a / (torch.abs(a) + 1)] | | | v | [检查 b.sum() < 0] | | | +-> 如果为真: | | | | | v | | [执行 __resume_at_30_2] | | | | | v | | [b = b * -1] | | | +-> 如果为假: | | | v | [执行 __resume_at_38_3] | +-> 如果 guard 条件不满足: | v [执行原始的 toy_example 函数] | v [计算 x = a / (torch.abs(a) + 1)] | v [检查 b.sum() < 0] | +-> 如果为真: | | | v | [b = b * -1] | +-> 如果为假: | v [不修改 b] | v
[返回 x * b] | v
结束这个流程图描述了 toy_example 函数的执行过程,包括了优化后的路径和原始路径。主要的区别在于是否满足 guard 条件,这决定了是执行优化后的代码还是原始代码。两种路径最终都会计算并返回 x * b。
然后可以问它__guard_0_for_torch_dynamo_resume_in_toy_example_at_8这个函数的作用：
这些检查都是根据输入Tensor的meta信息，python object信息，以及当前运行环境等来判断的。通过上面的流程图我们就可以一步步的去看torch.compile做了什么，例如__transformed_code_0_for_toy_example函数中__resume_at_30_2：
def __transformed_code_1_for_torch_dynamo_resume_in_toy_example_at_8(b, x): a = None # this line helps Python to generate bytecode with at least the same number of local variables as the original function __temp_9, = __compiled_fn_7(b, x) return __temp_9 # Note: if there is a transformed version below, this function might well not be executed directly. Please check the transformed version if possible.
def __resume_at_30_2(b, x): b = b * -1 return x * b def transformed___resume_at_30_2(b, x): __local_dict = {"b": b, "x": x} __global_dict = globals() if __guard_1_for_torch_dynamo_resume_in_toy_example_at_8(__local_dict, __global_dict): return __transformed_code_1_for_torch_dynamo_resume_in_toy_example_at_8(b, x) # Note: this function might well not be executed directly. It might well be transformed again, i.e. adding one more guards and transformed code. return __resume_at_30_2(b, x) def __transformed_code_0_for_toy_example(a, b): __temp_2, __temp_3 = __compiled_fn_1(a, b) x = __temp_2 if __temp_3: return __resume_at_30_2(b, x) return __resume_at_38_3(b, x)这个时候我们就轻松知道我们应该去查看__compiled_fn_7这个函数对应的编译产物了，如下图红色所示：
打开_compiled_fn_7_kernel0.py文件，我们可以看到原始的：
def __resume_at_30_2(b, x): b = b * -1 return x * b被fuse成了一个kernel，实现为：
cpp_fused_mul_0 = async_compile.cpp_pybinding(['const float*', 'const float*', 'float*'], '''
#include "/tmp/torchinductor_root/sk/cskh5dx62fglpphcrl6723dnmowdabouerrzy3dmqcngbxwfa7bv.h"
extern "C" void kernel(const float* in_ptr0, const float* in_ptr1, float* out_ptr0)
{ { #pragma omp simd simdlen(8) for(long x0=static_cast<long>(0L); x0<static_cast<long>(10L); x0+=static_cast<long>(1L)) { auto tmp0 = in_ptr0[static_cast<long>(x0)]; auto tmp1 = in_ptr1[static_cast<long>(x0)]; auto tmp2 = static_cast<float>(-1.0); auto tmp3 = decltype(tmp1)(tmp1 * tmp2); auto tmp4 = decltype(tmp0)(tmp0 * tmp3); out_ptr0[static_cast<long>(x0)] = tmp4; } }
}
''')对于cuda程序来说，整体流程也是类似的。
上面展示了一个完整的阅读depfy生成的torch.compile编译产物的例子，希望对大家有帮助。
六、使用PyTorch进行小样本学习的图像分类
近年来，基于深度学习的模型在目标检测和图像识别等任务中表现出色。像ImageNet这样具有挑战性的图像分类数据集，包含1000种不同的对象分类，现在一些模型已经超过了人类水平上。但是这些模型依赖于监督训练流程，标记训练数据的可用性对它们有重大影响，并且模型能够检测到的类别也仅限于它们接受训练的类。
由于在训练过程中没有足够的标记图像用于所有类，这些模型在现实环境中可能不太有用。并且我们希望的模型能够识别它在训练期间没有见到过的类，因为几乎不可能在所有潜在对象的图像上进行训练。我们将从几个样本中学习的问题被称为“少样本学习 Few-Shot learning”。
什么是小样本学习?
少样本学习是机器学习的一个子领域。它涉及到在只有少数训练样本和监督数据的情况下对新数据进行分类。只需少量的训练样本，我们创建的模型就可以相当好地执行。
考虑以下场景:在医疗领域，对于一些不常见的疾病，可能没有足够的x光图像用于训练。对于这样的场景，构建一个小样本学习分类器是完美的解决方案。
小样本的变化
一般来说，研究人员确定了四种类型:
- N-Shot Learning (NSL)
- Few-Shot Learning ( FSL )
- One-Shot Learning (OSL)
- Zero-Shot Learning (ZSL)
当我们谈论 FSL 时，我们通常指的是 N-way-K-Shot 分类。N 代表类别数，K 代表每个类中要训练的样本数。所以N-Shot Learning 被视为比所有其他概念更广泛的概念。可以说 Few-Shot、One-Shot 和 Zero-Shot是 NSL 的子领域。而零样本学习旨在在没有任何训练示例的情况下对看不见的类进行分类。
在 One-Shot Learning 中，每个类只有一个样本。Few-Shot 每个类有 2 到 5 个样本，也就是说 Few-Shot 是更灵活的 One-Shot Learning 版本。
小样本学习方法
通常，在解决 Few Shot Learning 问题时应考虑两种方法：
数据级方法 (DLA)
这个策略非常简单，如果没有足够的数据来创建实体模型并防止欠拟合和过拟合，那么就应该添加更多数据。正因为如此，许多 FSL 问题都可以通过利用来更大大的基础数据集的更多数据来解决。基本数据集的显着特征是它缺少构成我们对 Few-Shot 挑战的支持集的类。例如，如果我们想要对某种鸟类进行分类，则基础数据集可能包含许多其他鸟类的图片。
参数级方法 (PLA)
从参数级别的角度来看，Few-Shot Learning 样本相对容易过拟合，因为它们通常具有大的高维空间。限制参数空间、使用正则化和使用适当的损失函数将有助于解决这个问题。少量的训练样本将被模型泛化。
通过将模型引导到广阔的参数空间可以提高性能。由于缺乏训练数据，正常的优化方法可能无法产生准确的结果。
因为上面的原因，训练我们的模型以发现通过参数空间的最佳路径，产生最佳的预测结果。这种方法被称为元学习。
小样本学习图像分类算法
有4种比较常见的小样本学习的方法：
与模型无关的元学习 Model-Agnostic Meta-Learning
基于梯度的元学习 (GBML) 原则是 MAML 的基础。在 GBML 中，元学习者通过基础模型训练和学习所有任务表示的共享特征来获得先前的经验。每次有新任务要学习时，元学习器都会利用其现有经验和新任务提供的最少量的新训练数据进行微调训练。
一般情况下，如果我们随机初始化参数经过几次更新算法将不会收敛到良好的性能。MAML 试图解决这个问题。MAML 只需几个梯度步骤并且保证没有过度拟合的前提下，为元参数学习器提供了可靠的初始化，这样可以对新任务进行最佳快速学习。
步骤如下：
元学习者在每个分集（episode）开始时创建自己的副本C，
C 在这一分集上进行训练（在 base-model 的帮助下），
C 对查询集进行预测，
从这些预测中计算出的损失用于更新 C，
这种情况一直持续到完成所有分集的训练。
- 元学习者在每个分集（episode）开始时创建自己的副本C，
- C 在这一分集上进行训练（在 base-model 的帮助下），
- C 对查询集进行预测，
- 从这些预测中计算出的损失用于更新 C，
- 这种情况一直持续到完成所有分集的训练。
这种技术的最大优势在于，它被认为与元学习算法的选择无关。因此MAML 方法被广泛用于许多需要快速适应的机器学习算法，尤其是深度神经网。
匹配网络 Matching Networks
为解决 FSL 问题而创建的第一个度量学习方法是匹配网络 (MN)。
当使用匹配网络方法解决 Few-Shot Learning 问题时需要一个大的基础数据集。。
将该数据集分为几个分集之后，对于每一分集，匹配网络进行以下操作：
- 来自支持集和查询集的每个图像都被馈送到一个 CNN，该 CNN 为它们输出特征的嵌入
- 查询图像使用支持集训练的模型得到嵌入特征的余弦距离，通过 softmax 进行分类
- 分类结果的交叉熵损失通过 CNN 反向传播更新特征嵌入模型
匹配网络可以通过这种方式学习构建图像嵌入。MN 能够使用这种方法对照片进行分类，并且无需任何特殊的类别先验知识。他只要简单地比较类的几个实例就可以了。
由于类别因分集而异，因此匹配网络会计算对类别区分很重要的图片属性（特征）。而当使用标准分类时，算法会选择每个类别独有的特征。
原型网络 Prototypical Networks
与匹配网络类似的是原型网络（PN）。它通过一些细微的变化来提高算法的性能。PN 比 MN 取得了更好的结果，但它们训练过程本质上是相同的，只是比较了来自支持集的一些查询图片嵌入，但是 原型网络提供了不同的策略。
我们需要在 PN 中创建类的原型：通过对类中图像的嵌入进行平均而创建的类的嵌入。然后仅使用这些类原型来比较查询图像嵌入。当用于单样本学习问题时，它可与匹配网络相媲美。
关系网络 Relation Network
关系网络可以说继承了所有上面提到方法的研究的结果。RN是基于PN思想的但包含了显著的算法改进。
该方法使用的距离函数是可学习的，而不是像以前研究的事先定义它。关系模块位于嵌入模块之上，嵌入模块是从输入图像计算嵌入和类原型的部分。
可训练的关系模块（距离函数）输入是查询图像的嵌入与每个类的原型，输出为每个分类匹配的关系分数。关系分数通过 Softmax 得到一个预测。
使用 Open-AI Clip 进行零样本学习
CLIP（Contrastive Language-Image Pre-Training）是一个在各种（图像、文本）对上训练的神经网络。它无需直接针对任务进行优化，就可以为给定的图像来预测最相关的文本片段（类似于 GPT-2 和 3 的零样本的功能）。
CLIP 在 ImageNet“零样本”上可以达到原始 ResNet50 的性能，而且需要不使用任何标记示例，它克服了计算机视觉中的几个主要挑战，下面我们使用Pytorch来实现一个简单的分类模型。
引入包
! pip install ftfy regex tqdm
! pip install git+https://github.com/openai/CLIP.gitimport numpy as np
import torch
from pkg_resources import packagingprint("Torch version:", torch.__version__)加载模型
import clipclip.available\_models\(\) # it will list the names of available CLIP modelsmodel, preprocess = clip.load\("ViT-B/32"\) model.cuda\(\).eval\(\) input\_resolution = model.visual.input\_resolution context\_length = model.context\_length vocab\_size = model.vocab\_size print\("Model parameters:", f"\{np.sum\(\[int\(np.prod\(p.shape\)\) for p in model.parameters\(\)\]\):,\}"\) print\("Input resolution:", input\_resolution\) print\("Context length:", context\_length\) print\("Vocab size:", vocab\_size\)图像预处理
我们将向模型输入8个示例图像及其文本描述，并比较对应特征之间的相似性。
分词器不区分大小写，我们可以自由地给出任何合适的文本描述。
import os import skimage import IPython.display import matplotlib.pyplot as plt from PIL import Image import numpy as np from collections import OrderedDict import torch \%matplotlib inline \%config InlineBackend.figure\_format = 'retina' \# images in skimage to use and their textual descriptions descriptions = \{ "page": "a page of text about segmentation", "chelsea": "a facial photo of a tabby cat", "astronaut": "a portrait of an astronaut with the American flag", "rocket": "a rocket standing on a launchpad", "motorcycle\_right": "a red motorcycle standing in a garage", "camera": "a person looking at a camera on a tripod", "horse": "a black-and-white silhouette of a horse", "coffee": "a cup of coffee on a saucer" \}original\_images = \[\] images = \[\] texts = \[\] plt.figure\(figsize=\(16, 5\)\) for filename in \[filename for filename in os.listdir\(skimage.data\_dir\) if filename.endswith\(".png"\) or filename.endswith\(".jpg"\)\]: name = os.path.splitext\(filename\)\[0\] if name not in descriptions: continue image = Image.open\(os.path.join\(skimage.data\_dir, filename\)\).convert\("RGB"\) plt.subplot\(2, 4, len\(images\) + 1\) plt.imshow\(image\) plt.title\(f"\{filename\}\\n\{descriptions\[name\]\}"\) plt.xticks\(\[\]\) plt.yticks\(\[\]\) original\_images.append\(image\) images.append\(preprocess\(image\)\) texts.append\(descriptions\[name\]\) plt.tight\_layout\(\)结果的可视化如下：
我们对图像进行规范化，对每个文本输入进行标记，并运行模型的正传播获得图像和文本的特征。
image\_input = torch.tensor\(np.stack\(images\)\).cuda\(\) text\_tokens = clip.tokenize\(\["This is " + desc for desc in texts\]\).cuda\(\) with torch.no\_grad\(\): image\_features = model.encode\_image\(image\_input\).float\(\) text\_features = model.encode\_text\(text\_tokens\).float\(\)我们将特征归一化，并计算每一对的点积，进行余弦相似度计算
image\_features /= image\_features.norm\(dim=-1, keepdim=True\) text\_features /= text\_features.norm\(dim=-1, keepdim=True\) similarity = text\_features.cpu\(\).numpy\(\) \@ image\_features.cpu\(\).numpy\(\).T count = len\(descriptions\) plt.figure\(figsize=\(20, 14\)\) plt.imshow\(similarity, vmin=0.1, vmax=0.3\) \# plt.colorbar\(\) plt.yticks\(range\(count\), texts, fnotallow=18\) plt.xticks\(\[\]\) for i, image in enumerate\(original\_images\): plt.imshow\(image, extent=\(i - 0.5, i + 0.5, -1.6, -0.6\), origin="lower"\) for x in range\(similarity.shape\[1\]\): for y in range\(similarity.shape\[0\]\): plt.text\(x, y, f"\{similarity\[y, x\]:.2f\}", ha="center", va="center", size=12\) for side in \["left", "top", "right", "bottom"\]: plt.gca\(\).spines\[side\].set\_visible\(False\) plt.xlim\(\[-0.5, count - 0.5\]\) plt.ylim\(\[count + 0.5, -2\]\) plt.title\("Cosine similarity between text and image features", size=20\)
零样本的图像分类
from torchvision.datasets import CIFAR100 cifar100 = CIFAR100\(os.path.expanduser\("\~/.cache"\), transform=preprocess, download=True\) text\_descriptions = \[f"This is a photo of a \{label\}" for label in cifar100.classes\] text\_tokens = clip.tokenize\(text\_descriptions\).cuda\(\) with torch.no\_grad\(\): text\_features = model.encode\_text\(text\_tokens\).float\(\) text\_features /= text\_features.norm\(dim=-1, keepdim=True\) text\_probs = \(100.0 \* image\_features \@ text\_features.T\).softmax\(dim=-1\) top\_probs, top\_labels = text\_probs.cpu\(\).topk\(5, dim=-1\) plt.figure\(figsize=\(16, 16\)\) for i, image in enumerate\(original\_images\): plt.subplot\(4, 4, 2 \* i + 1\) plt.imshow\(image\) plt.axis\("off"\) plt.subplot\(4, 4, 2 \* i + 2\) y = np.arange\(top\_probs.shape\[-1\]\) plt.grid\(\) plt.barh\(y, top\_probs\[i\]\) plt.gca\(\).invert\_yaxis\(\) plt.gca\(\).set\_axisbelow\(True\) plt.yticks\(y, \[cifar100.classes\[index\] for index in top\_labels\[i\].numpy\(\)\]\) plt.xlabel\("probability"\) plt.subplots\_adjust\(wspace=0.5\) plt.show\(\)
可以看到，分类的效果还是非常好的
七、Python和C++中使用并行计算增强图像处理能力
你是否曾经发现自己要等待很长时间才能处理一个装满图片的文件夹？无论是将它们转换为灰度图还是执行其他图像处理任务，如果您在单个线程中处理所有内容，处理大型数据集可能会非常慢。幸运的是，并行计算可以解决这个问题！
本文中我们将介绍如何使用并行计算来加速常见的计算机视觉任务：将图像转换为灰度图。我们将研究两个示例：一个使用Joblib库在Python中编写，另一个使用 OpenMP 在 C++中编写。
Joblib简介：
OpenMP简介：
为什么要进行并行计算？
并行计算就像将一个大任务分成几个小任务，让多个工作人员同时处理每个任务。假设您要处理 100 张图像。与其同时处理一张图像，为什么不一次处理 5 张或 10 张呢？通过利用 CPU 的多个核心，可以节省大量时间。
🚀 Python与C++中的并行性
如果你使用 Python，你可能听说过全局解释器锁 (GIL)，它限制了线程的真正并行性。你可能想知道为什么 Python 在这方面会遇到困难。Python 的 GIL 确保单个进程中一次只有一个线程运行。这对于保证安全非常有用，但严重限制了图像处理等 CPU 密集型任务的性能。
另一方面，C++ 可以让你充分利用多线程的强大功能。你可以使用 OpenMP 轻松地将任务分配到不同的 CPU 内核上，以最小的努力实现真正的并行性。
🚀 Python中使用Joblib 加速
import os
import cv2
from glob import glob
from joblib import Parallel, delayedinput_folder = 'images/'
output_folder = 'output/'def convert_to_grayscale(image_path, output_folder):img = cv2.imread(image_path)gray_img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)output_path = os.path.join(output_folder, os.path.basename(image_path))cv2.imwrite(output_path, gray_img)if __name__ == "__main__":# get a list of abs path for all .jpg files whitin input_folderimage_files = glob(f"{input_folder}/*.jpg")# Parallel processing using JoblibParallel(n_jobs=-1, backend="threading")(delayed(convert_to_grayscale)(image, output_folder) for image in image_files)在这里，Joblib将图像处理任务拆分到多个 CPU 核心上。只需设置n_jobs=-1，所有可用核心都会被利用。
🚀 C++中使用OpenMP实现真正的多线程
OpenMP 是处理 C++ 中并行性的一种简单方法，可让您轻松并行化循环。对于将图像转换为灰度等 CPU 密集型任务，它非常高效。
首先，您需要安装 OpenCV 并设置 OpenMP。以下是快速设置指南：
1. 如果你还没有安装 OpenCV，先安装它
sudo apt-get install libopencv-dev2. 确保您拥有OpenMP（大多数现代编译器（如 GCC 和 Clang）都具有开箱即用的 OpenMP 支持）。
3. 要使用 OpenMP 编译 C++ 代码，请使用：
g++ -fopenmp -o image_converter image_converter.cpp `pkg-config --cflags --libs opencv4`fopenmp标志是启用OpenMP的标志。
现在，让我们看一下代码：
#include <opencv2/opencv.hpp>
#include <omp.h>
#include <filesystem>
#include <string>void convert_to_grayscale(const std::string& input_path, const std::string& output_folder)
{cv::Mat img = cv::imread(input_path);cv::Mat gray_img;cv::cvtColor(img, gray_img, cv::COLOR_BGR2GRAY);std::string output_path = output_folder + "/" + std::filesystem::path(input_path).filename().string();cv::imwrite(output_path, gray_img);
}int main()
{std::string input_folder = "images/";std::string output_folder = "output/";std::vector<std::string> image_files;for (const auto& entry : std::filesystem::directory_iterator(input_folder)){image_files.emplace_back(entry.path().string());}// Parallel processing using OpenMP#pragma omp parallel forfor (size_t i = 0; i < image_files.size(); ++i){convert_to_grayscale(image_files[i], output_folder);}return 0;
}代码中的#pragma omp parallel for让每个线程处理循环的一部分，将工作负载分布在多个 CPU 核心上。与Python不同，C++可以真正并行运行线程，而不会受到 GIL 的限制。
🚀 总结
通过利用并行计算，可以显著减少处理大量图像所需的时间。无论你使用 Python 还是 C++，都可以使用工具来加快工作流程。
在 Python 中，我们使用 Joblib 解决了 GIL 问题。对于更强大的多线程，C++ 与 OpenMP 的结合将带来翻天覆地的变化。
代码下载：
https://github.com/Gabriellgpc/multhreading-image-processing?source=post_page-----9a8895115ef8--------------------------------八、基于PyTorch FCOS模型实现螺丝检测
FCOS介绍
FCOS（Fully Convolutional One Stage）是于2019年由Zhi Tian等人提出的一种全卷积单阶段无锚框检测模型。该模型无需在训练时计算锚框的IoU值，极大的简化了模型的训练，使得训练方便，并且在推理时只有NMS（None Maximum Suppression，非及大值抑制），使得推理速度极快。在torchvision模型库中包含了以Resnet50作为骨干网络的FCOS，这样只需要创建模型进行训练即可。
模型介绍
FCOS模型的结构如图10.8所示，整个模型从左向右，可以分为三个部分：Backbone（骨干网络），Feature Pyramid（特征金字塔），以及由Classification、Center-ness和Regression构成的多尺度目标检测头作为输出。
Backbone主要用于提取提取不同尺度的特征，一般使用预训练的分类网络的特征提取部分作为BackBone，FCOS的Backbone，作为模型的入口，接受输入的图像，从Backbone中输出C3，C4和C5特征图。各特征图的高和宽标记在各特征图的左侧，分别为输入图像的1/8，1/16和1/32，对于输入为800×1024的图像，C3的尺寸为100×128，C4的尺寸为50×64，C5的尺寸为25×32。
Feature Pyramid接收来自骨干网络的C3，C4和C5特征图，一方面以C5特征图为基础构造尺寸为13×16的P6特征图和尺寸为7×8的P7特征图，另一方面将C5上采样和C4合并构造与C4同尺寸的P4，并将P4上采样和C3构造出与C4同尺寸的P3。这样通过Feature Pyramid总共构造出了P3，P4，P5，P6和P7共5个尺寸的特征图，目标检测就在这5个不同尺寸的特征图上完成。
多尺度目标检测头负责对从Feature Pyramid传来P3-7共5个特征图进行检测。在每个特征图上进行目标检测的头都具有相同的结构。每个头包含两个独立的部分，Classification结构和Center-ness结构共享一个部分，其中Classification结构输出大小为H×W×C形状的张量表示在H×W大小的特征图上检测总共C个类别的结果，Center-ness结构输出大小为H×W×1形状的张量表示在H×W的特征图上，是否为目标的中心；Regression独自为一个部分，输出大小为H×W×4形状的张量，表示以该元素为中心的目标到该元素的4个距离，如图10.9所示。图10.9显示了FCOS模型在Regression部分学习的4个距离，分别是中心到左边界的距离，中心到右边界的距离，中心到上边界的距离和中心到下边界的距离。这样FCOS模型就彻底抛弃了锚框IoU的计算，直接以这四个距离进行训练。
图10.8 FCOS目标检测模型
图10.9 FCOS目标外接矩形的表示
由于在torchvision中包含了以ResNet50为骨干网络的FCOS模型，并且还可以加载COCO上预训练的模型，这样创建一个FCOS模型就非常方便，创建方法与上一节创建预训练模型方法相同：
from torchvision.models import detection
model = detection.fcos_resnet50_fpn(progress=True,num_classes=3,pretrained_backbnotallow=True,trainable_backbone_layers=4)以上就可以创建一个具有检测3个类别，以带有预训练参数的ResNet50为骨干网络的FCOS模型，对于创建模型时其他参数及其含义可以参考API文档。
数据集制作
由于目标检测模型多样，因此，在训练前对于数据集的构建方法会有所差异。对于torchvision包中提供的所有目标检测模型已经对训练数据的格式进行了统一，因此，只需要把数据按照统一的方式进行构建后，torchvision包内的其它目标检测模型也可以使用。
由于通用目标检测数据集通常较大，不便于进行原理的演示。在这里使用一个样本量较小，类别数较小的目标检测数据集——螺丝螺母检测数据集。螺丝螺母检测数据集是一个开源目标检测数据集，下载地址为：
https://aistudio.baidu.com/aistudio/datasetdetail/6045同时，该数据集也附于本书的电子资料中。
图10.10 螺丝螺母检测数据集中的样本
螺丝螺母数据集包括413张训练集和10张测试集两部分。图10.10显示了一个带有标注的训练集中的样本，在样本中螺丝和螺母放置于一个白色托盘中，托盘放置在一灰色平台上，螺丝螺母使用矩形框进行标注。以下就以该数据集为例构造用于训练torchvision中模型的数据集。
使用torchvision中的目标检测模型训练自定义数据同样需要对数据集进行封装，并在得到样本的__getitem__()方法中返回一个表示样本元组的数据和标签，以(x, y)表示，其中x是一个范围为0-1的3×H×W的图像张量，y表示图像x的标签，是一个包含‘label’和‘boxes’两个键的字典，‘label’键里以整数张量的形式存储了图像中K个目标的标签值，‘boxes’键里存储了图像中K个对应目标外边矩形框的左上和右下共4个坐标值组成的一个K×4的数字张量，格式如下所示：
#样本标签y的格式：
{'labels': tensor([1, 1, 1, 2, 2, 2, 2, 2]),'boxes': tensor([[ 711, 233, 844, 506],[1036, 194, 1206, 459],[ 958, 406, 1239, 573],[1142, 194, 1275, 320],[ 780, 478, 908, 614],[ 766, 612, 914, 742],[ 972, 542, 1120, 678],[ 986, 684, 1120, 820]])}#以上表明样本中包含8个目标，3个目标的类别为1，5个的目标的类别为2。按照上述要求，通过继承torch.utils.data.Dataset类创建一个自定义的数据集，实现螺丝螺母数据集的构造：
from pathlib import Path
from torchvision.io import read_image,ImageReadMode
import json
import torch
class BNDataset(torch.utils.data.Dataset):def __init__(self, istrain=True,datapath='D:/data/lslm'): #注意修改数据集路径self.datadir=Path(datapath)/('train' if istrain else 'test')self.idxfile=self.datadir/('train.txt' if istrain else 'test.txt')self.labelnames=['background','bolt','nut']self.data=self.parseidxfile()def parseidxfile(self):lines=open(self.idxfile).readlines()return [line for line in lines if len(line)>5]def __getitem__(self, idx):data=self.data[idx].split('\t')x = read_image(str(self.datadir/data[0]),ImageReadMode.RGB)/255.0labels=[]boxes=[]for i in data[2:]:if len(i)<5: continuer=json.loads(i)labels.append(self.labelnames.index( r['value']))cords=r['coordinate']xyxy=cords[0][0],cords[0][1],cords[1][0],cords[1][1]boxes.append(xyxy)y = {'labels': torch.LongTensor(labels),'boxes': torch.tensor(boxes).long()}return x, ydef __len__(self):return len(self.data)以上代码对螺丝螺母数据集以BNDataset为类名进行封装，主要涉及的难点就是标签文件的解析，具体解析过程要结合上述代码和标签文件进行理解。在模型进行训练时，还需要使用把数据集进一步使用DataLoader封装：
def collate_fn(data):x = [i[0] for i in data]y = [i[1] for i in data]return x, ytrain_loader = torch.utils.data.DataLoader(dataset=BNDataset(istrain=True), batch_size=4, shuffle=True, drop_last=True, collate_fn=collate_fn)
test_loader = torch.utils.data.DataLoader(dataset=BNDataset(istrain=False), batch_size=1, shuffle=True, drop_last=True, collate_fn=collate_fn)在封装完成后，就可以使用上一节提到的可视化方法，进行样本和标签的可视化，以检查数据集构造的正确性，得到如图10.10所示的结果：
for i, (x, y) in enumerate(train_loader):labels=[loader.dataset.labelnames[i] for i in y[0]['labels']]
colors=[ ('red' if i=='nut' else 'blue') for i in labels]
image=draw_bounding_boxes(x[0], y[0]['boxes'],labels=labels,colors=colors,width=5,font_size=50,outtype='CHW')
vis.image(image)#图10.10所示
以上完成了螺丝螺母数据集的构造，能够用于FCOS模型的训练。下面介绍FCOS模型在该数据集上的训练以及模型的评估。
训练与预测
由于torchvision对FCOS模型进行了很好的封装，在准备好数据集后，训练方法与分类和分割网络的训练模式并无太大差异：创建优化器，构造损失函数，对数据集进行多次循环并根据反向传播的梯度进行参数的修正。将训练过程封装为train()函数，调用train()函数进行模型的训练，代码如下：
def train():model.train()optimizer = torch.optim.SGD(model.parameters(), lr=0.0001,momentum=0.98)for epoch in range(5):for i, (x, y) in enumerate(train_loader):outs = model(x, y)loss = outs['classification']+ outs['bbox_ctrness']+outs['bbox_regression']loss.backward()optimizer.step()optimizer.zero_grad()if i % 10 == 0:print(epoch, i, loss.item())torch.save(model, f'./models/tvs{epoch}.model')train()#输出结果：
0 0 2.1866644620895386
...
...
4 100 0.9854792356491089以上就是FCOS模型的训练代码，其中train()函数实现了模型的训练，在该函数中，将模型切换到训练模式，创建SGD优化器，总共训练5轮（可根据情况训练更多轮数），每10批打印损失值，经过4轮训练后，损失值从2.18降为了0.98，并且在每轮训练完成后都保存模型。
在模型训练完成后，就可以在测试集上查看和评估模型的检测效果。评估方法实质上与之前介绍的模型的使用方法是相同的，可以参考上一节的内容以便于理解。对模型在测试集上进行运行，并可视化结果，测试代码如下：
def test():model_load = torch.load('./models/tvs4.model')model_load.eval（)loader_test = torch.utils.data.DataLoader(dataset=BNDataset(istrain=False), batch_size=1, shuffle=False, drop_last=True, collate_fn=collate_fn)for i, (x, y) in enumerate(loader_test):with torch.no_grad():outs = model_load(x)res=outs[0]boxes=res['boxes']scores=res['scores']labels=res['labels']#阈值过滤threhold=0.5 #保留类别概率大于0.5的检测结果mask=scores>threholdscores=scores[mask]labels=labels[mask]boxes=boxes[mask]labelnames=[loader_test.dataset.labelnames[i]+f'{scores[idx]:.2f}' for idx,i in enumerate(labels)]colors=[ ('red' if i==1 else 'blue') for i in labels]img=draw_bounding_boxes(x[0],boxes,labels=labelnames,colors=colors,width=3,font_size=50,outtype=’CHW’)
vis.image(img)#使用visdom可视化，图10.11
图10.11显示了经过5轮的训练，FCOS模型在螺丝螺母数据集上的测试结果，其中（a），（b）中的螺丝和螺母均被正确的检出，特别是（a）中右侧两个螺母十分靠近也正确的得到检出，而（c），（d）中均有部分螺丝被错误检测。从结果可以看出，虽然仅仅进行了5轮的训练，但模型就已经能够较好地检出螺丝和螺母，可以预测经过更多轮数的训练能够取得更好的检测效果。
以上就是使用FCOS在螺丝螺母数据集上的训练和评估，对于torchvision中的其他模型的训练只需要修改模型，并在train()函数中根据模型的输出，修改计算损失各部分的构成和各部分损失的权重即可。
