Compare commits
2 Commits
3b7b64c915
...
b0030fe87f
Author | SHA1 | Date | |
---|---|---|---|
|
b0030fe87f | ||
|
260e3aa760 |
9
net.py
9
net.py
@ -7,6 +7,7 @@ from timm.models.layers import DropPath, to_2tuple, trunc_normal_
|
||||
from einops import rearrange
|
||||
|
||||
from componets.SCSA import SCSA
|
||||
from componets.WTConvCV2 import WTConv2d
|
||||
|
||||
|
||||
def drop_path(x, drop_prob: float = 0., training: bool = False):
|
||||
@ -248,6 +249,10 @@ class DetailNode(nn.Module):
|
||||
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)
|
||||
elif useBlock == 2:
|
||||
self.theta_phi = WTConv2d(in_channels=32, out_channels=32)
|
||||
self.theta_rho = WTConv2d(in_channels=32, out_channels=32)
|
||||
self.theta_eta = WTConv2d(in_channels=32, out_channels=32)
|
||||
else:
|
||||
self.theta_phi = InvertedResidualBlock(inp=32, oup=32, expand_ratio=2)
|
||||
self.theta_rho = InvertedResidualBlock(inp=32, oup=32, expand_ratio=2)
|
||||
@ -269,7 +274,7 @@ class DetailNode(nn.Module):
|
||||
class DetailFeatureFusion(nn.Module):
|
||||
def __init__(self, num_layers=3):
|
||||
super(DetailFeatureFusion, self).__init__()
|
||||
INNmodules = [DetailNode() for _ in range(num_layers)]
|
||||
INNmodules = [DetailNode(useBlock=2) for _ in range(num_layers)]
|
||||
self.net = nn.Sequential(*INNmodules)
|
||||
|
||||
def forward(self, x):
|
||||
@ -281,7 +286,7 @@ class DetailFeatureFusion(nn.Module):
|
||||
class DetailFeatureExtraction(nn.Module):
|
||||
def __init__(self, num_layers=3):
|
||||
super(DetailFeatureExtraction, self).__init__()
|
||||
INNmodules = [DetailNode() for _ in range(num_layers)]
|
||||
INNmodules = [DetailNode(useBlock=2) for _ in range(num_layers)]
|
||||
self.net = nn.Sequential(*INNmodules)
|
||||
|
||||
def forward(self, x):
|
||||
|
Loading…
Reference in New Issue
Block a user