Commit a6270e6d by 前钰

Upload New File

parent bcd56693
import torch # 导入 PyTorch 主库
import torch # 导入 PyTorch 主库
import torch.nn as nn # 导入神经网络相关模块
from mmcv.ops import DeformConv2d # 从 mmcv 库中导入 Deformable Convolution v2 模块
class DCNBlock(nn.Module): # 定义一个 DCNBlock 类,继承自 nn.Module
def __init__(self, in_channels, out_channels, kernel_size=3, padding=1):
super(DCNBlock, self).__init__() # 初始化父类
# 定义一个普通卷积用于学习 offset(形变卷积所需的位移参数)
self.offset_conv = nn.Conv2d(
in_channels, # 输入通道数
2 * kernel_size * kernel_size, # offset_conv 的输出通道为 2 * kernel_size * kernel_size 是因为:
# 对于每个卷积核位置,需要学习一个 (x_offset, y_offset) 的二维偏移量
kernel_size=3, # 使用3×3卷积提取offset
padding=1 # 保证输出大小一致
)
# 定义 Deformable Convolution(可变形卷积)层
self.dcn = DeformConv2d(
in_channels, # 输入通道数
out_channels, # 输出通道数
kernel_size=kernel_size, # 卷积核大小
padding=padding # 填充
)
# ReLU 激活函数
self.relu = nn.ReLU(inplace=True) # inplace=True节省内存
def forward(self, x): # 定义前向传播
offset = self.offset_conv(x) # 先通过普通卷积预测 offset(位移信息)
out = self.dcn(x, offset) # 将输入和 offset 一起传入可变形卷积
out = self.relu(out) # 输出通过 ReLU 激活
return out # 返回输出结果
# 测试代码
if __name__ == "__main__":
x = torch.randn(1, 64, 128, 128).cuda() # 构造一个 [1, 64, 128, 128] 的输入张量并放到 GPU
model = DCNBlock(in_channels=64, out_channels=64).cuda() # 创建 DCNBlock 模型并放到 GPU
out = model(x) # 进行前向传播
print("输入形状:", x.shape) # 打印输入张量的形状
print("输出形状:", out.shape) # 打印输出张量的形状
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment