| import torch |
| import torch.nn as nn |
|
|
| class PlainCNN(nn.Module): |
| def __init__(self, layer=12): |
| super().__init__() |
| self.head = nn.Sequential( |
| nn.Conv2d(3, 32, 3, 1, 1), nn.SiLU(), |
| nn.Conv2d(32, 64, 3, 1, 1), nn.SiLU(), |
| ) |
| layers = [] |
| for i in range(layer-4): |
| layers.append(nn.Conv2d(64, 64, 3, 1, 1)) |
| layers.append(nn.SiLU()) |
| self.middle_part = nn.ModuleList(layers) |
| self.tail = nn.Sequential( |
| nn.Conv2d(64, 32, 3, 1, 1), nn.SiLU(), |
| nn.Conv2d(32, 3, 3, 1, 1) |
| ) |
|
|
| def forward(self, x): |
| |
| res = x |
| x = self.head(x) |
| for b in self.middle_part: |
| x = b(x) |
| x = self.tail(x) |
| x = x + res |
| return x |
|
|
|
|
|
|
