|
马上注册,结交更多好友,享用更多功能^_^
您需要 登录 才可以下载或查看,没有账号?立即注册
x
class DeeperAutoencoderWithMoreAttention(nn.Module):
def __init__(self):
super(DeeperAutoencoderWithMoreAttention, self).__init__()
# Encoder
self.encoder = nn.Sequential(
ResidualBlockWithAttention(3, 32),
ResidualBlockWithAttention(32, 32),
nn.MaxPool2d(2, stride=2),
ResidualBlockWithAttention(32, 64),
ResidualBlockWithAttention(64, 64),
nn.MaxPool2d(2, stride=2),
ResidualBlockWithAttention(64, 128),
ResidualBlockWithAttention(128, 128),
nn.MaxPool2d(2, stride=2),
ResidualBlockWithAttention(128, 256),
ResidualBlockWithAttention(256, 256),
nn.MaxPool2d(2, stride=2),
ResidualBlockWithAttention(256, 512),
nn.MaxPool2d(2, stride=2),
)
# Decoder
self.decoder = nn.Sequential(
nn.Upsample(scale_factor=2),
ResidualBlockWithAttention(512, 512),
nn.Upsample(scale_factor=2),
ResidualBlockWithAttention(512, 256),
ResidualBlockWithAttention(256, 256),
nn.Upsample(scale_factor=2),
ResidualBlockWithAttention(256, 128),
ResidualBlockWithAttention(128, 128),
nn.Upsample(scale_factor=2),
ResidualBlockWithAttention(128, 64),
ResidualBlockWithAttention(64, 64),
nn.Upsample(scale_factor=2),
ResidualBlockWithAttention(64, 32),
ResidualBlockWithAttention(32, 32),
nn.Conv2d(32, 3, kernel_size=3, stride=1, padding=1),
nn.ReLU()
)
def forward(self, x):
x = self.encoder(x)
x = self.decoder(x)
return x |
|