|
import logging |
|
|
|
import torch.nn as nn |
|
|
|
from siclib.models.base_model import BaseModel |
|
from siclib.models.utils.modules import ConvModule |
|
|
|
logger = logging.getLogger(__name__) |
|
|
|
|
|
|
|
|
|
|
|
class LowLevelEncoder(BaseModel): |
|
default_conf = { |
|
"feat_dim": 64, |
|
"in_channel": 3, |
|
"keep_resolution": True, |
|
} |
|
|
|
required_data_keys = ["image"] |
|
|
|
def _init(self, conf): |
|
logger.debug(f"Initializing LowLevelEncoder with {conf}") |
|
|
|
if self.conf.keep_resolution: |
|
self.conv1 = ConvModule(conf.in_channel, conf.feat_dim, kernel_size=3, padding=1) |
|
self.conv2 = ConvModule(conf.feat_dim, conf.feat_dim, kernel_size=3, padding=1) |
|
else: |
|
self.conv1 = nn.Conv2d( |
|
conf.in_channel, conf.feat_dim, kernel_size=7, stride=2, padding=3, bias=False |
|
) |
|
self.bn1 = nn.BatchNorm2d(conf.feat_dim) |
|
self.relu = nn.ReLU(inplace=True) |
|
|
|
def _forward(self, data): |
|
x = data["image"] |
|
|
|
assert ( |
|
x.shape[-1] % 32 == 0 and x.shape[-2] % 32 == 0 |
|
), "Image size must be multiple of 32 if not using single image input." |
|
|
|
if self.conf.keep_resolution: |
|
c1 = self.conv1(x) |
|
c2 = self.conv2(c1) |
|
else: |
|
x = self.conv1(x) |
|
x = self.bn1(x) |
|
c2 = self.relu(x) |
|
|
|
return {"features": c2} |
|
|
|
def loss(self, pred, data): |
|
raise NotImplementedError |
|
|
