program

Signed-off-by: 葛峻恺 <202115006@mail.sdu.edu.cn>

Network/Model.py

@@ -0,0 +1,137 @@
+import torch
+import torch.nn as nn
+
+# Define 1D convolution with kernel size 5
+def conv1d_5(inplanes, outplanes, stride=1):
+    return nn.Conv1d(inplanes, outplanes, kernel_size=5, stride=stride,
+                     padding=2, bias=False)
+
+# Transformer-based self-attention module
+class TransformerLayer(nn.Module):
+    def __init__(self, embed_dim, num_heads, ff_dim, dropout=0.1):
+        super(TransformerLayer, self).__init__()
+        self.attention = nn.MultiheadAttention(embed_dim, num_heads, dropout=dropout, batch_first=True)
+        self.norm1 = nn.LayerNorm(embed_dim)
+        self.ffn = nn.Sequential(
+            nn.Linear(embed_dim, ff_dim),
+            nn.ReLU(),
+            nn.Linear(ff_dim, embed_dim)
+        )
+        self.norm2 = nn.LayerNorm(embed_dim)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x):
+        x = x.permute(0, 2, 1)  # [B, C, L] -> [B, L, C]
+        attn_output, _ = self.attention(x, x, x)
+        x = self.norm1(x + self.dropout(attn_output))
+        
+        ffn_output = self.ffn(x)
+        x = self.norm2(x + self.dropout(ffn_output))
+        
+        x = x.permute(0, 2, 1)  # [B, L, C] -> [B, C, L]
+        return x
+
+# Downsampling Block using ResNet-style skip connections
+class Block(nn.Module):
+    def __init__(self, inplanes, planes, stride=1, downsample=None, bn=False):
+        super(Block, self).__init__()
+        self.bn = bn
+        self.conv1 = conv1d_5(inplanes, planes, stride)
+        self.bn1 = nn.BatchNorm1d(planes)
+        self.relu = nn.ReLU(inplace=False)
+        self.conv2 = conv1d_5(planes, planes)
+        self.bn2 = nn.BatchNorm1d(planes)
+        self.downsample = downsample
+
+    def forward(self, x):
+        residual = x
+        out = self.conv1(x)
+        if self.bn:
+            out = self.bn1(out)
+        out = self.relu(out)
+        
+        out = self.conv2(out)
+        if self.bn:
+            out = self.bn2(out)
+        out = self.relu(out)
+        
+        if self.downsample is not None:
+            residual = self.downsample(x)
+        
+        out = out + residual
+        out = self.relu(out)
+        return out
+
+# Upsampling Block
+class Decoder_block(nn.Module):
+    def __init__(self, inplanes, outplanes, kernel_size=5, stride=5):
+        super(Decoder_block, self).__init__()
+        self.upsample = nn.ConvTranspose1d(inplanes, outplanes,
+                                           kernel_size=kernel_size, stride=stride, bias=False)
+        self.conv1 = conv1d_5(inplanes, outplanes)
+        self.relu = nn.ReLU(inplace=False)
+        self.conv2 = conv1d_5(outplanes, outplanes)
+
+    def forward(self, x1, x2):
+        x1 = self.upsample(x1)
+        out = torch.cat((x1, x2), dim=1)
+        out = self.conv1(out)
+        out = self.relu(out)
+        out = self.conv2(out)
+        out = self.relu(out)
+        return out
+
+# Main Model
+class Model(nn.Module):
+    def __init__(self, inplanes=1, outplanes=2, layers=[2, 2, 2, 2]):
+        super(Model, self).__init__()
+        self.inplanes = inplanes
+        self.outplanes = outplanes
+        self.encoder1 = self._make_encoder(Block, 32, layers[0], 5)
+        self.encoder2 = self._make_encoder(Block, 64, layers[1], 5)
+        self.encoder3 = self._make_encoder(Block, 128, layers[2], 5)
+        self.encoder4 = self._make_encoder(Block, 256, layers[3], 4)
+        
+        # Self-Attention Layer between Encoder and Decoder
+        self.self_attention = TransformerLayer(embed_dim=256, num_heads=8, ff_dim=512)
+        
+        self.decoder3 = Decoder_block(256, 128, stride=4, kernel_size=4)
+        self.decoder2 = Decoder_block(128, 64)
+        self.decoder1 = Decoder_block(64, 32)
+        self.conv1x1 = nn.ConvTranspose1d(32, outplanes, kernel_size=5, stride=5, bias=False)
+
+    def _make_encoder(self, block, planes, blocks, stride=1):
+        downsample = None
+        if self.inplanes != planes or stride != 1:
+            downsample = nn.Conv1d(self.inplanes, planes, kernel_size=1, stride=stride, bias=False)
+        layers = [block(self.inplanes, planes, stride, downsample)]
+        self.inplanes = planes
+        for _ in range(1, blocks):
+            layers.append(block(self.inplanes, planes))
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        down1 = self.encoder1(x)
+        down2 = self.encoder2(down1)
+        down3 = self.encoder3(down2)
+        down4 = self.encoder4(down3)
+        
+        # Apply self-attention layer
+        attention_out = self.self_attention(down4)
+        
+        up3 = self.decoder3(attention_out, down3)
+        up2 = self.decoder2(up3, down2)
+        up1 = self.decoder1(up2, down1)
+        out = self.conv1x1(up1)
+        return out
+
+# Test function to verify input-output compatibility
+if __name__ == "__main__":
+    model = Model(inplanes=1, outplanes=1, layers=[3, 3, 3, 3])
+    model.eval()
+    image = torch.randn(1, 1, 1000)
+    with torch.no_grad():
+        output = model(image)
+    print(output.size())
+
+