vztu / maxim-pytorch

[CVPR 2022 Oral] PyTorch re-implementation for "MAXIM: Multi-Axis MLP for Image Processing", with *training code*. Official Jax repo: https://github.com/google-research/maxim

Hi there, I am Zhengzhong Tu 👋

我正在招生，详情请见个人主页.

I'm an Assistant Professor of Computer Science at Texas A&M University. Previously, I was an AI researcher at Google Research.

I graduated from UT Austin ECE, where I worked with Professor Alan Bovik in the LIVE Lab.

📫 Contact me by:

• Website
• 知乎:「假熊猫」
• Twitter
• Linkedin
• Scholar

maxim-pytorch's Issues

Where is the maxim? /maxim-pytorch/basicsr/models/archs/restormer?

Hello？ Is there really a maxim model here? What I found is a restomer model. If you can, please explain this part. Please.

from basicsr.models import create_model in train.py

Hello dear...
is there a create_model in basicsr ?!

training.yml for LOL enhancement

Thank you for your excellent work about image filed firstly? In this repo， only Deblurring,Denoising,Deraining can be seen . Could you provide image enhancement training.py and .yml code? Looking forward to your reply.

Multiscale MAXIM

Hello,
First of all, thank you so much for sharing the maxim architecture in PyTorch. Unfortunately, the provided code does not work with a num_supervision_scales higher than 1. I have changed the code slightly but I have not been successful to solve the issues yet. For example, in line 1115, we have buttleneckblocks, while according to the paper we need backbone blocks (depth 4 in Table 9 of the paper). Where are these blocks?
According to issue #3, I have downloaded the weights, but the name of layers are different from your code. Therefore, the PyTorch model cannot load the weights properly.
Finally, according to the paper maxim has around 14 million parameters for image enhancement, but the PyTorch model has around 400 million parameters. What is the problem? I used the following code to calculate the number of parameters.
print(sum(p.numel() for p in maxim.parameters() if p.requires_grad))
Can you please let me know when the final version of the model will be available?

When will the pytorch version be released?

Code not finished yet, plz stay tuned. Will close this issue until mostly done

file size changed after jax2torch

Why did I change the size from 260MB to 105MB after converting the pre training weight to torch weight? Is this normal? Is it because the Torch framework's volume optimization is done better?

How is the patch input into the network?

This network occupies too much graphics memory. Usually, inputs of 4, 1, 128, and 128 will occupy up to 16GB of graphics memory. High definition images must be input in patches. How did the original JAX network handle this?

will the pretrained model be open sourced ?

Excuse me, are the GridGatingUnit and other modules not completed?

Excuse me, are the GridGatingUnit and other modules not completed?

How to load pretrained Enhancement model in maxim_torch.py

I downloaded weights from maxim repo
Then in maxim-pytorch repo i run jax2torch.py file with

python maxim_pytorch/jax2torch.py -c maxim_ckpt_Enhancement_FiveK_checkpoint.npz
It works and I get torch_weight.pth file

I then try to load it but im unable to understand if Im giving the wrong arguments or your code is wrong

from maxim_pytorch.maxim_torch import MAXIM_dns_3s
import torch
import cv2
import numpy as np
from torchvision import transforms

from pathlib import Path


# These params are from https://github.com/google-research/maxim/blob/3c8265171ffccc80c3c9124844aef0d381609956/maxim/models/maxim.py#L910
s2 = {
    "features": 32,
    "depth": 3,
    "num_stages": 2, #
    "num_groups": 2, # 
    "num_bottleneck_blocks": 2, #
    "block_gmlp_factor": 2,
    "grid_gmlp_factor": 2,
    "input_proj_factor": 2,
    "channels_reduction": 4,
}

model = MAXIM_dns_3s(features=32, depth=3, block_gmlp_factor=2, grid_gmlp_factor=2, input_proj_factor=2, channels_reduction=4, num_supervision_scales=2)
state = torch.load("torch_weight.pth")

model.load_state_dict(state)
model.eval()

I get error:

RuntimeError: Error(s) in loading state_dict for MAXIM_dns_3s:
	Unexpected key(s) in state_dict: "stage_1_output_conv_0.bias", "stage_1_output_conv_0.weight", "stage_1_output_conv_1.bias", "stage_1_output_conv_1.weight", "stage_1_output_conv_2.bias", "stage_1_output_conv_2.weight".