有没有NonLocal BLOCK的模块分享

发表于 2025-4-28 13:12:24

有没有NonLocal BLOCK的模块分享或者现在最新的NonLocal模块；或者在哪里import的

发表于 2025-4-28 13:12:25

在Python机器学习库PyTorch中，你可以自己编写一个NonLocal Block，下面是一个可能的实现示例：

import torch
from torch import nn
class NonLocalBlock(nn.Module):
def __init__(self, in_channels, inter_channels=None):
super(NonLocalBlock, self).__init__()
self.in_channels = in_channels
self.inter_channels = inter_channels
if self.inter_channels is None:
self.inter_channels = in_channels // 2
if self.inter_channels == 0:
self.inter_channels = 1
self.g = nn.Conv2d(in_channels=self.in_channels, out_channels=self.inter_channels, kernel_size=1, stride=1, padding=0)
self.W = nn.Sequential(
nn.Conv2d(in_channels=self.inter_channels, out_channels=self.in_channels, kernel_size=1, stride=1, padding=0),
nn.BatchNorm2d(self.in_channels)
)
self.W[0].weight.data.fill_(0)
self.W[1].weight.data.fill_(0)
self.W[1].bias.data.fill_(0)
self.theta = nn.Conv2d(in_channels=self.in_channels, out_channels=self.inter_channels, kernel_size=1, stride=1, padding=0)
self.phi = nn.Conv2d(in_channels=self.in_channels, out_channels=self.inter_channels, kernel_size=1, stride=1, padding=0)
def forward(self, x):
batch_size = x.size(0)
g_x = self.g(x).view(batch_size, self.inter_channels, -1)
g_x = g_x.permute(0, 2, 1)
theta_x = self.theta(x).view(batch_size, self.inter_channels, -1)
theta_x = theta_x.permute(0, 2, 1)
phi_x = self.phi(x).view(batch_size, self.inter_channels, -1)
f = torch.matmul(theta_x, phi_x)
f_div_C = nn.functional.softmax(f, dim=-1)
y = torch.matmul(f_div_C, g_x)
y = y.permute(0, 2, 1).contiguous()
y = y.view(batch_size, self.inter_channels, *x.size()[2:])
W_y = self.W(y)
z = W_y + x
return z

复制代码

以上定义的`NonLocalBlock`模块是基于2D卷积的版本，适用于图像和视频等输入。

想要使用这个NonLocal Block，你只需要将其作为一个模块添加到你的网络架构中即可。例如，如果你在构建一个CNN，并想在两个卷积层之间添加一个`NonLocalBlock`，你可以这样做：

class MyCNN(nn.Module):
def __init__(self, num_classes):
super(MyCNN, self).__init__()
self.Conv1 = nn.Conv2d(in_channels=3, out_channels=64, kernel_size=3, stride=1, padding=1)
self.NonLocalBlock = NonLocalBlock(in_channels=64)
self.Conv2 = nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, stride=1, padding=1)
def forward(self, x):
x = self.Conv1(x)
x = self.NonLocalBlock(x)
x = self.Conv2(x)
return x

复制代码

这种Self-attention机制的模块可以捕获图像中远距离的依赖关系，对于某些任务(如语义分割)可能非常有用。

[已完结] 有没有NonLocal BLOCK的模块分享