Shouldn't this rather be
shifted_x = torch.roll(x, shifts=(-shift_pl, -shift_lat, -shift_lon), dims=(1, 2, 3))

the un-shifting does implement asymmetric shifting correctly it seems like.

For the pangu scripts, there doesn't appear to be any modifications made to the data to normalise it. I can see that the 'cal_mean_std.py' file calculates the mean & standard deviation of each variable but this isn't then used anywhere else from what I can tell.

Is this intentional?

Hi,
I am trying to test FuXi() model on a weather dataset. However, I am facing an error when running it through the SwinTransformerV2Block(). The details are:

The input dimension is (B, in_chans = 1, T = 2, 46, 90)
patch_size=(2, 4, 4)
window_size = 8

While passing through the UTransformer block -> self.layer(x) -> SwinTransformerV2Block.forward() -> window_partition function, I am facing a dimension mismatch at this line..

Would you please help? It looks like there is an issue with the padding function, or the window_size is not compatible.

My code:

B = 1  
in_chans = 1
out_chans = 1  # number of input channels or output channels
input = torch.randn(B, in_chans, 2, 46, 90)  # B C T Lat Lon
model = fuxi.Fuxi(img_size=(2, 46, 90), 
                patch_size=(2, 4, 4), 
                in_chans=in_chans, out_chans=out_chans,
                embed_dim=128, 
                num_groups=32, 
                num_heads=4, 
                window_size=8)  

output = model(input)  # B C Lat Lon

Error:

46 x = x.view(B, H // window_size[0], window_size[0], W // window_size[1], window_size[1], C)
47 windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size[0], window_size[1], C)
48 return windows

RuntimeError: shape '[1, 1, 8, 2, 8, 128]' is invalid for input of size 18432

Your support would be very helpful.

Hi, first of all, thank you, this is super useful.
I was wondering about the skip connection between layer 1 and layer 4 of Pangu Weather model.
In your implementation the skip is added after layer4, just before patch recovery, following original Pangu pseudo-code https://github.com/198808xc/Pangu-Weather/blob/403b7b69ee66f220f78afeeafa117dd0f52c0b7d/pseudocode.py#L236

However, this does not match the figure of PanguWeather (seems to indicate that skip is before layer4), and intuitively adding the skip right before patch recovery is way too late to do any meaningful processing of the skip connection. I think it might be an error in their pseudo-code.

Thoughts ?

Hi Lizhuoq, thanks for this code - it has been really helpful. The Panguweather portion of this code is an accurate implementation of the official Panguweather repository in Pytorch. That is to say, it takes the full Pangu 24 hr .onnx model and runs a prediction + the code to train a pangu_lite & run this locally.

I'm looking to change the format of Panguweather slightly for downscaling and wanted to get your thoughts on taking the full .onnx model and finetuning the first few layers of this rather than training a completely new model. In examples\pangu-lite\train.py, how would I amend this to take up where the full Pangu .onnx model left off?

Appreciate any suggestions

Hello!
I have a number of questions about launching FuXi:

1. How much is estimated it may take to launch the model (32GB of CUDA was not enough for me to test run the model from your examples (There may be a memory leak?))?
2. The original model has an input parameter encoding the time shift ("temb" from "time_encoding"). Is there any way to combine the weights of the model and retrain this network from Pytorch ?

P.S. My task is fine-tuning the model on a one-hour time grid.