feat(net): 为 DetailNode模块添加可选卷积块
- 在 DetailNode 类中引入 useBlock 参数,用于选择不同的卷积块 - 新增 DepthwiseSeparableConvBlock 类,实现深度可分离卷积 - 根据 useBlock 的值,选择使用 DepthwiseSeparableConvBlock 或 InvertedResidualBlock - 优化了网络结构,提供了更多的灵活性和选择性
This commit is contained in:
parent
a435e035c0
commit
0336fc23ba
26
net.py
26
net.py
@ -222,11 +222,33 @@ class InvertedResidualBlock(nn.Module):
|
|||||||
def forward(self, x):
|
def forward(self, x):
|
||||||
return self.bottleneckBlock(x)
|
return self.bottleneckBlock(x)
|
||||||
|
|
||||||
|
class DepthwiseSeparableConvBlock(nn.Module):
|
||||||
|
def __init__(self, inp, oup, kernel_size=3, stride=1, padding=1):
|
||||||
|
super(DepthwiseSeparableConvBlock, self).__init__()
|
||||||
|
self.depthwise = nn.Conv2d(inp, inp, kernel_size, stride, padding, groups=inp, bias=False)
|
||||||
|
self.pointwise = nn.Conv2d(inp, oup, 1, bias=False)
|
||||||
|
self.bn = nn.BatchNorm2d(oup)
|
||||||
|
self.relu = nn.ReLU(inplace=True)
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
x = self.depthwise(x)
|
||||||
|
x = self.pointwise(x)
|
||||||
|
x = self.bn(x)
|
||||||
|
x = self.relu(x)
|
||||||
|
return x
|
||||||
|
|
||||||
class DetailNode(nn.Module):
|
class DetailNode(nn.Module):
|
||||||
def __init__(self):
|
def __init__(self,useBlock=0):
|
||||||
super(DetailNode, self).__init__()
|
super(DetailNode, self).__init__()
|
||||||
|
if useBlock == 0:
|
||||||
|
self.theta_phi = DepthwiseSeparableConvBlock(inp=32, oup=32)
|
||||||
|
self.theta_rho = DepthwiseSeparableConvBlock(inp=32, oup=32)
|
||||||
|
self.theta_eta = DepthwiseSeparableConvBlock(inp=32, oup=32)
|
||||||
|
elif useBlock == 1:
|
||||||
|
self.theta_phi = InvertedResidualBlock(inp=32, oup=32, expand_ratio=2)
|
||||||
|
self.theta_rho = InvertedResidualBlock(inp=32, oup=32, expand_ratio=2)
|
||||||
|
self.theta_eta = InvertedResidualBlock(inp=32, oup=32, expand_ratio=2)
|
||||||
|
else:
|
||||||
self.theta_phi = InvertedResidualBlock(inp=32, oup=32, expand_ratio=2)
|
self.theta_phi = InvertedResidualBlock(inp=32, oup=32, expand_ratio=2)
|
||||||
self.theta_rho = InvertedResidualBlock(inp=32, oup=32, expand_ratio=2)
|
self.theta_rho = InvertedResidualBlock(inp=32, oup=32, expand_ratio=2)
|
||||||
self.theta_eta = InvertedResidualBlock(inp=32, oup=32, expand_ratio=2)
|
self.theta_eta = InvertedResidualBlock(inp=32, oup=32, expand_ratio=2)
|
||||||
|
Loading…
Reference in New Issue
Block a user