whaifree/pfcfuse
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Update net.py

Browse Source
This commit is contained in:
HXY13 2024-06-29 10:54:33 +08:00 committed by GitHub
parent 282d38e52a
commit f0765f0835
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 0 additions and 20 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

20
net.py
View File

@ -1,4 +1,3 @@
# poolformer
import torch import torch
import torch.nn as nn import torch.nn as nn
import math import math
@ -9,14 +8,6 @@ from einops import rearrange
def drop_path(x, drop_prob: float = 0., training: bool = False): 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: if drop_prob == 0. or not training:
return x return x
keep_prob = 1 - drop_prob keep_prob = 1 - drop_prob
@ -41,7 +32,6 @@ class DropPath(nn.Module):
def forward(self, x): def forward(self, x):
return drop_path(x, self.drop_prob, self.training) return drop_path(x, self.drop_prob, self.training)
class Pooling(nn.Module): class Pooling(nn.Module):
def __init__(self, kernel_size=3): def __init__(self, kernel_size=3):
super().__init__() super().__init__()
@ -72,13 +62,6 @@ class PoolMlp(nn.Module):
self.act = act_layer() self.act = act_layer()
self.fc2 = nn.Conv2d(hidden_features, out_features, 1, bias=bias) self.fc2 = nn.Conv2d(hidden_features, out_features, 1, bias=bias)
self.drop = nn.Dropout(drop) 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): def forward(self, x):
x = self.fc1(x) # (B, C, H, W) --> (B, C, H, W) x = self.fc1(x) # (B, C, H, W) --> (B, C, H, W)
@ -88,7 +71,6 @@ class PoolMlp(nn.Module):
x = self.drop(x) x = self.drop(x)
return x return x
class BaseFeatureExtraction(nn.Module): class BaseFeatureExtraction(nn.Module):
def __init__(self, dim, pool_size=3, mlp_ratio=4., def __init__(self, dim, pool_size=3, mlp_ratio=4.,
act_layer=nn.GELU, act_layer=nn.GELU,
@ -369,7 +351,6 @@ class Restormer_Encoder(nn.Module):
self.encoder_level1 = nn.Sequential( self.encoder_level1 = nn.Sequential(
*[TransformerBlock(dim=dim, num_heads=heads[0], ffn_expansion_factor=ffn_expansion_factor, *[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])]) bias=bias, LayerNorm_type=LayerNorm_type) for i in range(num_blocks[0])])
self.baseFeature = BaseFeatureExtraction(dim=dim) self.baseFeature = BaseFeatureExtraction(dim=dim)
self.detailFeature = DetailFeatureExtraction() self.detailFeature = DetailFeatureExtraction()
@ -424,4 +405,3 @@ if __name__ == '__main__':
window_size = 8 window_size = 8
modelE = Restormer_Encoder().cuda() modelE = Restormer_Encoder().cuda()
modelD = Restormer_Decoder().cuda() modelD = Restormer_Decoder().cuda()