Thank you very much for your outstanding work. I am a little curious. In my inference test, the memory consumption of your model is much lower than that of the transformer-based method. However, this is not the case in your paper. I wonder if there is something wrong with my method. The code is as follows. I look forward to your reply. Thanks.
import time
from typing import Tuple, List
import numpy as np
import tqdm
from torch import Tensor, optim
import torch
import torch.nn as nn
import torch.nn.functional as F
from einops.layers.torch import Rearrange
from basicsr.utils.registry import ARCH_REGISTRY
######################
# Meta Architecture
######################
@ARCH_REGISTRY.register()
class SeemoRe(nn.Module):
def __init__(self,
scale: int = 4,
in_chans: int = 3,
num_experts: int = 6,
num_layers: int = 6,
embedding_dim: int = 64,
img_range: float = 1.0,
use_shuffle: bool = False,
global_kernel_size: int = 11,
recursive: int = 2,
lr_space: int = 1,
topk: int = 2, ):
super().__init__()
self.scale = scale
self.num_in_channels = in_chans
self.num_out_channels = in_chans
self.img_range = img_range
rgb_mean = (0.4488, 0.4371, 0.4040)
self.mean = torch.Tensor(rgb_mean).view(1, 3, 1, 1)
# -- SHALLOW FEATURES --
self.conv_1 = nn.Conv2d(self.num_in_channels, embedding_dim, kernel_size=3, padding=1)
# -- DEEP FEATURES --
self.body = nn.ModuleList(
[ResGroup(in_ch=embedding_dim,
num_experts=num_experts,
use_shuffle=use_shuffle,
topk=topk,
lr_space=lr_space,
recursive=recursive,
global_kernel_size=global_kernel_size) for i in range(num_layers)]
)
# -- UPSCALE --
self.norm = LayerNorm(embedding_dim, data_format='channels_first')
self.conv_2 = nn.Conv2d(embedding_dim, embedding_dim, kernel_size=3, padding=1)
self.upsampler = nn.Sequential(
nn.Conv2d(embedding_dim, (scale ** 2) * self.num_out_channels, kernel_size=3, padding=1),
nn.PixelShuffle(scale)
)
def forward(self, x: torch.Tensor) -> torch.Tensor:
self.mean = self.mean.type_as(x)
x = (x - self.mean) * self.img_range
# -- SHALLOW FEATURES --
x = self.conv_1(x)
res = x
# -- DEEP FEATURES --
for idx, layer in enumerate(self.body):
x = layer(x)
x = self.norm(x)
# -- HR IMAGE RECONSTRUCTION --
x = self.conv_2(x) + res
x = self.upsampler(x)
x = x / self.img_range + self.mean
return x
#############################
# Components
#############################
class ResGroup(nn.Module):
def __init__(self,
in_ch: int,
num_experts: int,
global_kernel_size: int = 11,
lr_space: int = 1,
topk: int = 2,
recursive: int = 2,
use_shuffle: bool = False):
super().__init__()
self.local_block = RME(in_ch=in_ch,
num_experts=num_experts,
use_shuffle=use_shuffle,
lr_space=lr_space,
topk=topk,
recursive=recursive)
self.global_block = SME(in_ch=in_ch,
kernel_size=global_kernel_size)
def forward(self, x: torch.Tensor) -> torch.Tensor:
x = self.local_block(x)
x = self.global_block(x)
return x
#############################
# Global Block
#############################
class SME(nn.Module):
def __init__(self,
in_ch: int,
kernel_size: int = 11):
super().__init__()
self.norm_1 = LayerNorm(in_ch, data_format='channels_first')
self.block = StripedConvFormer(in_ch=in_ch, kernel_size=kernel_size)
self.norm_2 = LayerNorm(in_ch, data_format='channels_first')
self.ffn = GatedFFN(in_ch, mlp_ratio=2, kernel_size=3, act_layer=nn.GELU())
def forward(self, x: torch.Tensor) -> torch.Tensor:
x = self.block(self.norm_1(x)) + x
x = self.ffn(self.norm_2(x)) + x
return x
class StripedConvFormer(nn.Module):
def __init__(self,
in_ch: int,
kernel_size: int):
super().__init__()
self.in_ch = in_ch
self.kernel_size = kernel_size
self.padding = kernel_size // 2
self.proj = nn.Conv2d(in_ch, in_ch, kernel_size=1, padding=0)
self.to_qv = nn.Sequential(
nn.Conv2d(in_ch, in_ch * 2, kernel_size=1, padding=0),
nn.GELU(),
)
self.attn = StripedConv2d(in_ch, kernel_size=kernel_size, depthwise=True)
def forward(self, x: torch.Tensor) -> torch.Tensor:
q, v = self.to_qv(x).chunk(2, dim=1)
q = self.attn(q)
x = self.proj(q * v)
return x
#############################
# Local Blocks
#############################
class RME(nn.Module):
def __init__(self,
in_ch: int,
num_experts: int,
topk: int,
lr_space: int = 1,
recursive: int = 2,
use_shuffle: bool = False, ):
super().__init__()
self.norm_1 = LayerNorm(in_ch, data_format='channels_first')
self.block = MoEBlock(in_ch=in_ch, num_experts=num_experts, topk=topk, use_shuffle=use_shuffle,
recursive=recursive, lr_space=lr_space, )
self.norm_2 = LayerNorm(in_ch, data_format='channels_first')
self.ffn = GatedFFN(in_ch, mlp_ratio=2, kernel_size=3, act_layer=nn.GELU())
def forward(self, x: torch.Tensor) -> torch.Tensor:
x = self.block(self.norm_1(x)) + x
x = self.ffn(self.norm_2(x)) + x
return x
#################
# MoE Layer
#################
class MoEBlock(nn.Module):
def __init__(self,
in_ch: int,
num_experts: int,
topk: int,
use_shuffle: bool = False,
lr_space: str = "linear",
recursive: int = 2):
super().__init__()
self.use_shuffle = use_shuffle
self.recursive = recursive
self.conv_1 = nn.Sequential(
nn.Conv2d(in_ch, in_ch, kernel_size=3, padding=1),
nn.GELU(),
nn.Conv2d(in_ch, 2 * in_ch, kernel_size=1, padding=0)
)
self.agg_conv = nn.Sequential(
nn.Conv2d(in_ch, in_ch, kernel_size=4, stride=4, groups=in_ch),
nn.GELU())
self.conv = nn.Sequential(
nn.Conv2d(in_ch, in_ch, kernel_size=3, stride=1, padding=1, groups=in_ch),
nn.Conv2d(in_ch, in_ch, kernel_size=1, padding=0)
)
self.conv_2 = nn.Sequential(
StripedConv2d(in_ch, kernel_size=3, depthwise=True),
nn.GELU())
if lr_space == "linear":
grow_func = lambda i: i + 2
elif lr_space == "exp":
grow_func = lambda i: 2 ** (i + 1)
elif lr_space == "double":
grow_func = lambda i: 2 * i + 2
else:
raise NotImplementedError(f"lr_space {lr_space} not implemented")
self.moe_layer = MoELayer(
experts=[Expert(in_ch=in_ch, low_dim=grow_func(i)) for i in range(num_experts)],
# add here multiple of 2 as low_dim
gate=Router(in_ch=in_ch, num_experts=num_experts),
num_expert=topk,
)
self.proj = nn.Conv2d(in_ch, in_ch, kernel_size=1, padding=0)
def calibrate(self, x: torch.Tensor) -> torch.Tensor:
b, c, h, w = x.shape
res = x
for _ in range(self.recursive):
x = self.agg_conv(x)
x = self.conv(x)
x = F.interpolate(x, size=(h, w), mode="bilinear", align_corners=False)
return res + x
def forward(self, x: torch.Tensor) -> torch.Tensor:
x = self.conv_1(x)
if self.use_shuffle:
x = channel_shuffle(x, groups=2)
x, k = torch.chunk(x, chunks=2, dim=1)
x = self.conv_2(x)
k = self.calibrate(k)
x = self.moe_layer(x, k)
x = self.proj(x)
return x
class MoELayer(nn.Module):
def __init__(self, experts: List[nn.Module], gate: nn.Module, num_expert: int = 1):
super().__init__()
assert len(experts) > 0
self.experts = nn.ModuleList(experts)
self.gate = gate
self.num_expert = num_expert
def forward(self, inputs: torch.Tensor, k: torch.Tensor):
out = self.gate(inputs)
weights = F.softmax(out, dim=1, dtype=torch.float).to(inputs.dtype)
topk_weights, topk_experts = torch.topk(weights, self.num_expert)
out = inputs.clone()
if self.training:
exp_weights = torch.zeros_like(weights)
exp_weights.scatter_(1, topk_experts, weights.gather(1, topk_experts))
for i, expert in enumerate(self.experts):
out += expert(inputs, k) * exp_weights[:, i:i + 1, None, None]
else:
selected_experts = [self.experts[i] for i in topk_experts.squeeze(dim=0)]
for i, expert in enumerate(selected_experts):
out += expert(inputs, k) * topk_weights[:, i:i + 1, None, None]
return out
class Expert(nn.Module):
def __init__(self,
in_ch: int,
low_dim: int, ):
super().__init__()
self.conv_1 = nn.Conv2d(in_ch, low_dim, kernel_size=1, padding=0)
self.conv_2 = nn.Conv2d(in_ch, low_dim, kernel_size=1, padding=0)
self.conv_3 = nn.Conv2d(low_dim, in_ch, kernel_size=1, padding=0)
def forward(self, x: torch.Tensor, k: torch.Tensor) -> torch.Tensor:
x = self.conv_1(x)
x = self.conv_2(k) * x # here no more sigmoid
x = self.conv_3(x)
return x
class Router(nn.Module):
def __init__(self,
in_ch: int,
num_experts: int):
super().__init__()
self.body = nn.Sequential(
nn.AdaptiveAvgPool2d(1),
Rearrange('b c 1 1 -> b c'),
nn.Linear(in_ch, num_experts, bias=False),
)
def forward(self, x: torch.Tensor) -> torch.Tensor:
return self.body(x)
#################
# Utilities
#################
class StripedConv2d(nn.Module):
def __init__(self,
in_ch: int,
kernel_size: int,
depthwise: bool = False):
super().__init__()
self.in_ch = in_ch
self.kernel_size = kernel_size
self.padding = kernel_size // 2
self.conv = nn.Sequential(
nn.Conv2d(in_ch, in_ch, kernel_size=(1, self.kernel_size), padding=(0, self.padding),
groups=in_ch if depthwise else 1),
nn.Conv2d(in_ch, in_ch, kernel_size=(self.kernel_size, 1), padding=(self.padding, 0),
groups=in_ch if depthwise else 1),
)
def forward(self, x: torch.Tensor) -> torch.Tensor:
return self.conv(x)
def channel_shuffle(x, groups=2):
bat_size, channels, w, h = x.shape
group_c = channels // groups
x = x.view(bat_size, groups, group_c, w, h)
x = torch.transpose(x, 1, 2).contiguous()
x = x.view(bat_size, -1, w, h)
return x
class GatedFFN(nn.Module):
def __init__(self,
in_ch,
mlp_ratio,
kernel_size,
act_layer, ):
super().__init__()
mlp_ch = in_ch * mlp_ratio
self.fn_1 = nn.Sequential(
nn.Conv2d(in_ch, mlp_ch, kernel_size=1, padding=0),
act_layer,
)
self.fn_2 = nn.Sequential(
nn.Conv2d(in_ch, in_ch, kernel_size=1, padding=0),
act_layer,
)
self.gate = nn.Conv2d(mlp_ch // 2, mlp_ch // 2,
kernel_size=kernel_size, padding=kernel_size // 2, groups=mlp_ch // 2)
def feat_decompose(self, x):
s = x - self.gate(x)
x = x + self.sigma * s
return x
def forward(self, x: torch.Tensor):
x = self.fn_1(x)
x, gate = torch.chunk(x, 2, dim=1)
gate = self.gate(gate)
x = x * gate
x = self.fn_2(x)
return x
class LayerNorm(nn.Module):
r""" LayerNorm that supports two data formats: channels_last (default) or channels_first.
The ordering of the dimensions in the inputs. channels_last corresponds to inputs with
shape (batch_size, height, width, channels) while channels_first corresponds to inputs
with shape (batch_size, channels, height, width).
"""
def __init__(self, normalized_shape, eps=1e-6, data_format="channels_last"):
super().__init__()
self.weight = nn.Parameter(torch.ones(normalized_shape))
self.bias = nn.Parameter(torch.zeros(normalized_shape))
self.eps = eps
self.data_format = data_format
if self.data_format not in ["channels_last", "channels_first"]:
raise NotImplementedError
self.normalized_shape = (normalized_shape,)
def forward(self, x):
if self.data_format == "channels_last":
return F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
elif self.data_format == "channels_first":
u = x.mean(1, keepdim=True)
s = (x - u).pow(2).mean(1, keepdim=True)
x = (x - u) / torch.sqrt(s + self.eps)
x = self.weight[:, None, None] * x + self.bias[:, None, None]
return x
def build(upscale=4):
return SeemoRe(scale=upscale,
in_chans=3,
num_experts=3,
num_layers=16,
embedding_dim=48,
img_range=1.0,
use_shuffle=False,
global_kernel_size=11,
recursive=1,
lr_space="exp",
topk=1, )
def main():
# Set device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# Create model instance
model = build(4).to(device).eval()
# Generate random input data
input_shape = (1, 3, 320, 180) # Batch size of 1
input_data = torch.randn(input_shape).to(device)
# Set number of inferences
num_inferences = 1000
# Warm up
print('Warming up...\n')
with torch.inference_mode():
for _ in range(5):
_ = model(input_data)
# Reset CUDA memory stats
torch.cuda.reset_peak_memory_stats()
# Initialize CUDA events
starter, ender = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
timings = np.zeros((num_inferences, 1))
# Measure inference time
print('Measuring inference time...\n')
with torch.inference_mode():
for i in tqdm.tqdm(range(num_inferences)):
starter.record()
output = model(input_data)
ender.record()
torch.cuda.synchronize() # Wait for the events to be recorded
timings[i] = starter.elapsed_time(ender) # Time in milliseconds
average_inference_time = np.mean(timings)
# Calculate memory usage
memory_allocated = torch.cuda.memory_allocated(device) / (1024 ** 2) # MB
max_memory_allocated = torch.cuda.max_memory_allocated(device) / (1024 ** 2) # MB
max_memory_reserved = torch.cuda.max_memory_reserved(device) / (1024 ** 2) # MB
# Calculate parameters and MACs using thop
macs, params = profile(model, inputs=(input_data,))
# Output results
print(f"Input shape: {input_shape}")
print(f"Output shape: {output.shape}")
print(f"Average inference time over {num_inferences} runs: {average_inference_time:.4f} ms")
print(f"Memory allocated: {memory_allocated:.2f} MB")
print(f"Max memory allocated: {max_memory_allocated:.2f} MB")
print(f"Max memory reserved: {max_memory_reserved:.2f} MB")
print(f"Number of parameters: {params / 1e3:.2f} K")
print(f"MACs: {macs / 1e9:.2f} G")
if __name__ == '__main__':
from thop import profile
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
main()
`
Thank you very much for your outstanding work. I am a little curious. In my inference test, the memory consumption of your model is much lower than that of the transformer-based method. However, this is not the case in your paper. I wonder if there is something wrong with my method. The code is as follows. I look forward to your reply. Thanks.
