{"id":15601046,"url":"https://github.com/lucidrains/metnet3-pytorch","last_synced_at":"2025-04-06T08:11:54.956Z","repository":{"id":205481180,"uuid":"714337467","full_name":"lucidrains/metnet3-pytorch","owner":"lucidrains","description":"Implementation of MetNet-3, SOTA neural weather model out of Google Deepmind, in Pytorch","archived":false,"fork":false,"pushed_at":"2023-11-16T16:50:23.000Z","size":1114,"stargazers_count":211,"open_issues_count":4,"forks_count":26,"subscribers_count":6,"default_branch":"main","last_synced_at":"2025-03-30T07:07:56.149Z","etag":null,"topics":["artificial-intelligence","deep-learning","unet","vision-transformers","weather-forecasting"],"latest_commit_sha":null,"homepage":"","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"mit","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/lucidrains.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null}},"created_at":"2023-11-04T16:01:51.000Z","updated_at":"2025-03-03T06:51:58.000Z","dependencies_parsed_at":null,"dependency_job_id":"dd2aed4c-2659-4646-8ee7-a7e53f6b97cc","html_url":"https://github.com/lucidrains/metnet3-pytorch","commit_stats":{"total_commits":32,"total_committers":1,"mean_commits":32.0,"dds":0.0,"last_synced_commit":"a0b107d7f4b792f612341a33b63b54f79200f998"},"previous_names":["lucidrains/metnet-3"],"tags_count":12,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lucidrains%2Fmetnet3-pytorch","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lucidrains%2Fmetnet3-pytorch/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lucidrains%2Fmetnet3-pytorch/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lucidrains%2Fmetnet3-pytorch/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/lucidrains","download_url":"https://codeload.github.com/lucidrains/metnet3-pytorch/tar.gz/refs/heads/main","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":247451654,"owners_count":20940939,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["artificial-intelligence","deep-learning","unet","vision-transformers","weather-forecasting"],"created_at":"2024-10-03T02:12:58.273Z","updated_at":"2025-04-06T08:11:54.937Z","avatar_url":"https://github.com/lucidrains.png","language":"Python","funding_links":[],"categories":["Python"],"sub_categories":[],"readme":"\u003cimg src=\"./metnet3.png\" width=\"450px\"\u003e\u003c/img\u003e\n\n## MetNet-3 - Pytorch\n\nImplementation of \u003ca href=\"https://blog.research.google/2023/11/metnet-3-state-of-art-neural-weather.html\"\u003eMetNet 3\u003c/a\u003e, SOTA neural weather model out of Google Deepmind, in Pytorch\n\nThe model architecture is pretty unremarkable. It is basically a U-net with a specific \u003ca href=\"https://arxiv.org/abs/2204.01697\"\u003ewell performing vision transformer\u003c/a\u003e. The most interesting thing about the paper may end up being the loss scaling in section 4.3.2\n\n## Appreciation\n\n- \u003ca href=\"https://stability.ai/\"\u003eStabilityAI\u003c/a\u003e, \u003ca href=\"https://a16z.com/supporting-the-open-source-ai-community/\"\u003eA16Z Open Source AI Grant Program\u003c/a\u003e, and \u003ca href=\"https://huggingface.co/\"\u003e🤗 Huggingface\u003c/a\u003e for the generous sponsorships, as well as my other sponsors, for affording me the independence to open source current artificial intelligence research\n\n## Install\n\n```bash\n$ pip install metnet3-pytorch\n```\n\n## Usage\n\n```python\nimport torch\nfrom metnet3_pytorch import MetNet3\n\nmetnet3 = MetNet3(\n    dim = 512,\n    num_lead_times = 722,\n    lead_time_embed_dim = 32,\n    input_spatial_size = 624,\n    attn_dim_head = 8,\n    hrrr_channels = 617,\n    input_2496_channels = 2 + 14 + 1 + 2 + 20,\n    input_4996_channels = 16 + 1,\n    precipitation_target_bins = dict(\n        mrms_rate = 512,\n        mrms_accumulation = 512,\n    ),\n    surface_target_bins = dict(\n        omo_temperature = 256,\n        omo_dew_point = 256,\n        omo_wind_speed = 256,\n        omo_wind_component_x = 256,\n        omo_wind_component_y = 256,\n        omo_wind_direction = 180\n    ),\n    hrrr_loss_weight = 10,\n    hrrr_norm_strategy = 'sync_batchnorm',  # this would use a sync batchnorm to normalize the input hrrr and target, without having to precalculate the mean and variance of the hrrr dataset per channel\n    hrrr_norm_statistics = None             # you can also also set `hrrr_norm_strategy = \"precalculated\"` and pass in the mean and variance as shape `(2, 617)` through this keyword argument\n)\n\n# inputs\n\nlead_times = torch.randint(0, 722, (2,))\nhrrr_input_2496 = torch.randn((2, 617, 624, 624))\nhrrr_stale_state = torch.randn((2, 1, 624, 624))\ninput_2496 = torch.randn((2, 39, 624, 624))\ninput_4996 = torch.randn((2, 17, 624, 624))\n\n# targets\n\nprecipitation_targets = dict(\n    mrms_rate = torch.randint(0, 512, (2, 512, 512)),\n    mrms_accumulation = torch.randint(0, 512, (2, 512, 512)),\n)\n\nsurface_targets = dict(\n    omo_temperature = torch.randint(0, 256, (2, 128, 128)),\n    omo_dew_point = torch.randint(0, 256, (2, 128, 128)),\n    omo_wind_speed = torch.randint(0, 256, (2, 128, 128)),\n    omo_wind_component_x = torch.randint(0, 256, (2, 128, 128)),\n    omo_wind_component_y = torch.randint(0, 256, (2, 128, 128)),\n    omo_wind_direction = torch.randint(0, 180, (2, 128, 128))\n)\n\nhrrr_target = torch.randn(2, 617, 128, 128)\n\ntotal_loss, loss_breakdown = metnet3(\n    lead_times = lead_times,\n    hrrr_input_2496 = hrrr_input_2496,\n    hrrr_stale_state = hrrr_stale_state,\n    input_2496 = input_2496,\n    input_4996 = input_4996,\n    precipitation_targets = precipitation_targets,\n    surface_targets = surface_targets,\n    hrrr_target = hrrr_target\n)\n\ntotal_loss.backward()\n\n# after much training from above, you can predict as follows\n\nmetnet3.eval()\n\nsurface_preds, hrrr_pred, precipitation_preds = metnet3(\n    lead_times = lead_times,\n    hrrr_input_2496 = hrrr_input_2496,\n    hrrr_stale_state = hrrr_stale_state,\n    input_2496 = input_2496,\n    input_4996 = input_4996,\n)\n\n\n# Dict[str, Tensor], Tensor, Dict[str, Tensor]\n```\n\n## Todo\n\n- [x] figure out all the cross entropy and MSE losses\n- [x] auto-handle normalization across all the channels of the HRRR by tracking a running mean and variance of targets during training (using sync batchnorm as hack)\n- [x] allow researcher to pass in their own normalization variables for HRRR\n- [x] build all the inputs to spec, also make sure hrrr input is normalized, offer option to unnormalize hrrr predictions\n- [x] make sure model can be easily saved and loaded, with different ways of handling hrrr norm\n\n- [ ] figure out the topological embedding, consult a neural weather researcher\n\n## Citations\n\n```bibtex\n@article{Andrychowicz2023DeepLF,\n    title   = {Deep Learning for Day Forecasts from Sparse Observations},\n    author  = {Marcin Andrychowicz and Lasse Espeholt and Di Li and Samier Merchant and Alexander Merose and Fred Zyda and Shreya Agrawal and Nal Kalchbrenner},\n    journal = {ArXiv},\n    year    = {2023},\n    volume  = {abs/2306.06079},\n    url     = {https://api.semanticscholar.org/CorpusID:259129311}\n}\n```\n\n```bibtex\n@inproceedings{ElNouby2021XCiTCI,\n    title   = {XCiT: Cross-Covariance Image Transformers},\n    author  = {Alaaeldin El-Nouby and Hugo Touvron and Mathilde Caron and Piotr Bojanowski and Matthijs Douze and Armand Joulin and Ivan Laptev and Natalia Neverova and Gabriel Synnaeve and Jakob Verbeek and Herv{\\'e} J{\\'e}gou},\n    booktitle = {Neural Information Processing Systems},\n    year    = {2021},\n    url     = {https://api.semanticscholar.org/CorpusID:235458262}\n}\n```\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Flucidrains%2Fmetnet3-pytorch","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Flucidrains%2Fmetnet3-pytorch","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Flucidrains%2Fmetnet3-pytorch/lists"}