diff --git a/net.py b/net.py index a359de9..e404ed0 100644 --- a/net.py +++ b/net.py @@ -1,4 +1,3 @@ -# poolformer import torch import torch.nn as nn import math @@ -9,14 +8,6 @@ from einops import rearrange def drop_path(x, drop_prob: float = 0., training: bool = False): - """ - Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks). - This is the same as the DropConnect impl I created for EfficientNet, etc networks, however, - the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper... - See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for - changing the layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use - 'survival rate' as the argument. - """ if drop_prob == 0. or not training: return x keep_prob = 1 - drop_prob @@ -41,7 +32,6 @@ class DropPath(nn.Module): def forward(self, x): return drop_path(x, self.drop_prob, self.training) - class Pooling(nn.Module): def __init__(self, kernel_size=3): super().__init__() @@ -72,13 +62,6 @@ class PoolMlp(nn.Module): self.act = act_layer() self.fc2 = nn.Conv2d(hidden_features, out_features, 1, bias=bias) self.drop = nn.Dropout(drop) - # self.apply(self._init_weights) - - # def _init_weights(self, m): - # if isinstance(m, nn.Conv2D): - # trunc_normal_(m.weight) - # if m.bias is not None: - # zeros_(m.bias) def forward(self, x): x = self.fc1(x) # (B, C, H, W) --> (B, C, H, W) @@ -88,7 +71,6 @@ class PoolMlp(nn.Module): x = self.drop(x) return x - class BaseFeatureExtraction(nn.Module): def __init__(self, dim, pool_size=3, mlp_ratio=4., act_layer=nn.GELU, @@ -369,7 +351,6 @@ class Restormer_Encoder(nn.Module): self.encoder_level1 = nn.Sequential( *[TransformerBlock(dim=dim, num_heads=heads[0], ffn_expansion_factor=ffn_expansion_factor, bias=bias, LayerNorm_type=LayerNorm_type) for i in range(num_blocks[0])]) - self.baseFeature = BaseFeatureExtraction(dim=dim) self.detailFeature = DetailFeatureExtraction() @@ -424,4 +405,3 @@ if __name__ == '__main__': window_size = 8 modelE = Restormer_Encoder().cuda() modelD = Restormer_Decoder().cuda() -