We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Re-think how configs are handled in train.py
Currently configuration to `train.py` is handled with [OmegaConf](https://omegaconf.readthedocs.io/en/2.1_branch/usage.html). This made more sense when the tasks (and accompanying trainer code) were fragmented, as we could easily define per-task configuration. Now that the trainer code that we would like to include in base TorchGeo are being generalized into things like `ClassificationTask` and `SemanticSegmentationTask` _and_ it is clear that more complicated training configurations won't be supported by torchgeo proper, it might make sense to pull out the OmegaConf part, and go with a more simple `argparse` based approach. Bonus: this would also allow us to get rid of a dependency. I'm not sure how exactly the argparse approach would work in all cases but it is worth more thought!

Lightning has a few pieces of docs that can help with this:
- https://pytorch-lightning.readthedocs.io/en/stable/common/trainer.html#trainer-in-python-scripts
- https://pytorch-lightning.readthedocs.io/en/stable/common/hyperparameters.html
- https://pytorch-lightning.readthedocs.io/en/latest/common/lightning_cli.html

Whatever we settle on here should definitely still allow passing arguments via a YAML config file. This allows reproducible benchmark experiment configurations to be saved in source control.

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `train.py`
Content:
```
1 #!/usr/bin/env python3
2 
3 # Copyright (c) Microsoft Corporation. All rights reserved.
4 # Licensed under the MIT License.
5 
6 """torchgeo model training script."""
7 
8 import os
9 from typing import Any, Dict, Tuple, Type, cast
10 
11 import pytorch_lightning as pl
12 from omegaconf import DictConfig, OmegaConf
13 from pytorch_lightning import loggers as pl_loggers
14 from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint
15 
16 from torchgeo.datamodules import (
17     BigEarthNetDataModule,
18     ChesapeakeCVPRDataModule,
19     COWCCountingDataModule,
20     CycloneDataModule,
21     ETCI2021DataModule,
22     EuroSATDataModule,
23     LandCoverAIDataModule,
24     NAIPChesapeakeDataModule,
25     OSCDDataModule,
26     RESISC45DataModule,
27     SEN12MSDataModule,
28     So2SatDataModule,
29     UCMercedDataModule,
30 )
31 from torchgeo.trainers import (
32     BYOLTask,
33     ClassificationTask,
34     MultiLabelClassificationTask,
35     RegressionTask,
36     SemanticSegmentationTask,
37 )
38 
39 TASK_TO_MODULES_MAPPING: Dict[
40     str, Tuple[Type[pl.LightningModule], Type[pl.LightningDataModule]]
41 ] = {
42     "bigearthnet_all": (MultiLabelClassificationTask, BigEarthNetDataModule),
43     "bigearthnet_s1": (MultiLabelClassificationTask, BigEarthNetDataModule),
44     "bigearthnet_s2": (MultiLabelClassificationTask, BigEarthNetDataModule),
45     "byol": (BYOLTask, ChesapeakeCVPRDataModule),
46     "chesapeake_cvpr_5": (SemanticSegmentationTask, ChesapeakeCVPRDataModule),
47     "chesapeake_cvpr_7": (SemanticSegmentationTask, ChesapeakeCVPRDataModule),
48     "chesapeake_cvpr_prior": (SemanticSegmentationTask, ChesapeakeCVPRDataModule),
49     "cowc_counting": (RegressionTask, COWCCountingDataModule),
50     "cyclone": (RegressionTask, CycloneDataModule),
51     "eurosat": (ClassificationTask, EuroSATDataModule),
52     "etci2021": (SemanticSegmentationTask, ETCI2021DataModule),
53     "landcoverai": (SemanticSegmentationTask, LandCoverAIDataModule),
54     "naipchesapeake": (SemanticSegmentationTask, NAIPChesapeakeDataModule),
55     "oscd_all": (SemanticSegmentationTask, OSCDDataModule),
56     "oscd_rgb": (SemanticSegmentationTask, OSCDDataModule),
57     "resisc45": (ClassificationTask, RESISC45DataModule),
58     "sen12ms_all": (SemanticSegmentationTask, SEN12MSDataModule),
59     "sen12ms_s1": (SemanticSegmentationTask, SEN12MSDataModule),
60     "sen12ms_s2_all": (SemanticSegmentationTask, SEN12MSDataModule),
61     "sen12ms_s2_reduced": (SemanticSegmentationTask, SEN12MSDataModule),
62     "so2sat_supervised": (ClassificationTask, So2SatDataModule),
63     "so2sat_unsupervised": (ClassificationTask, So2SatDataModule),
64     "ucmerced": (ClassificationTask, UCMercedDataModule),
65 }
66 
67 
68 def set_up_omegaconf() -> DictConfig:
69     """Loads program arguments from either YAML config files or command line arguments.
70 
71     This method loads defaults/a schema from "conf/defaults.yaml" as well as potential
72     arguments from the command line. If one of the command line arguments is
73     "config_file", then we additionally read arguments from that YAML file. One of the
74     config file based arguments or command line arguments must specify task.name. The
75     task.name value is used to grab a task specific defaults from its respective
76     trainer. The final configuration is given as merge(task_defaults, defaults,
77     config file, command line). The merge() works from the first argument to the last,
78     replacing existing values with newer values. Additionally, if any values are
79     merged into task_defaults without matching types, then there will be a runtime
80     error.
81 
82     Returns:
83         an OmegaConf DictConfig containing all the validated program arguments
84 
85     Raises:
86         FileNotFoundError: when ``config_file`` does not exist
87         ValueError: when ``task.name`` is not a valid task
88     """
89     conf = OmegaConf.load("conf/defaults.yaml")
90     command_line_conf = OmegaConf.from_cli()
91 
92     if "config_file" in command_line_conf:
93         config_fn = command_line_conf.config_file
94         if not os.path.isfile(config_fn):
95             raise FileNotFoundError(f"config_file={config_fn} is not a valid file")
96 
97         user_conf = OmegaConf.load(config_fn)
98         conf = OmegaConf.merge(conf, user_conf)
99 
100     conf = OmegaConf.merge(  # Merge in any arguments passed via the command line
101         conf, command_line_conf
102     )
103 
104     # These OmegaConf structured configs enforce a schema at runtime, see:
105     # https://omegaconf.readthedocs.io/en/2.0_branch/structured_config.html#merging-with-other-configs
106     task_name = conf.experiment.task
107     task_config_fn = os.path.join("conf", "task_defaults", f"{task_name}.yaml")
108     if task_name == "test":
109         task_conf = OmegaConf.create()
110     elif os.path.exists(task_config_fn):
111         task_conf = cast(DictConfig, OmegaConf.load(task_config_fn))
112     else:
113         raise ValueError(
114             f"experiment.task={task_name} is not recognized as a valid task"
115         )
116 
117     conf = OmegaConf.merge(task_conf, conf)
118     conf = cast(DictConfig, conf)  # convince mypy that everything is alright
119 
120     return conf
121 
122 
123 def main(conf: DictConfig) -> None:
124     """Main training loop."""
125     ######################################
126     # Setup output directory
127     ######################################
128 
129     experiment_name = conf.experiment.name
130     task_name = conf.experiment.task
131     if os.path.isfile(conf.program.output_dir):
132         raise NotADirectoryError("`program.output_dir` must be a directory")
133     os.makedirs(conf.program.output_dir, exist_ok=True)
134 
135     experiment_dir = os.path.join(conf.program.output_dir, experiment_name)
136     os.makedirs(experiment_dir, exist_ok=True)
137 
138     if len(os.listdir(experiment_dir)) > 0:
139         if conf.program.overwrite:
140             print(
141                 f"WARNING! The experiment directory, {experiment_dir}, already exists, "
142                 + "we might overwrite data in it!"
143             )
144         else:
145             raise FileExistsError(
146                 f"The experiment directory, {experiment_dir}, already exists and isn't "
147                 + "empty. We don't want to overwrite any existing results, exiting..."
148             )
149 
150     with open(os.path.join(experiment_dir, "experiment_config.yaml"), "w") as f:
151         OmegaConf.save(config=conf, f=f)
152 
153     ######################################
154     # Choose task to run based on arguments or configuration
155     ######################################
156     # Convert the DictConfig into a dictionary so that we can pass as kwargs.
157     task_args = cast(Dict[str, Any], OmegaConf.to_object(conf.experiment.module))
158     datamodule_args = cast(
159         Dict[str, Any], OmegaConf.to_object(conf.experiment.datamodule)
160     )
161 
162     datamodule: pl.LightningDataModule
163     task: pl.LightningModule
164     if task_name in TASK_TO_MODULES_MAPPING:
165         task_class, datamodule_class = TASK_TO_MODULES_MAPPING[task_name]
166         task = task_class(**task_args)
167         datamodule = datamodule_class(**datamodule_args)
168     else:
169         raise ValueError(
170             f"experiment.task={task_name} is not recognized as a valid task"
171         )
172 
173     ######################################
174     # Setup trainer
175     ######################################
176     tb_logger = pl_loggers.TensorBoardLogger(conf.program.log_dir, name=experiment_name)
177 
178     checkpoint_callback = ModelCheckpoint(
179         monitor="val_loss", dirpath=experiment_dir, save_top_k=1, save_last=True
180     )
181     early_stopping_callback = EarlyStopping(
182         monitor="val_loss", min_delta=0.00, patience=18
183     )
184 
185     trainer_args = cast(Dict[str, Any], OmegaConf.to_object(conf.trainer))
186 
187     trainer_args["callbacks"] = [checkpoint_callback, early_stopping_callback]
188     trainer_args["logger"] = tb_logger
189     trainer_args["default_root_dir"] = experiment_dir
190     trainer = pl.Trainer(**trainer_args)
191 
192     if trainer_args.get("auto_lr_find"):
193         trainer.tune(model=task, datamodule=datamodule)
194 
195     ######################################
196     # Run experiment
197     ######################################
198     trainer.fit(model=task, datamodule=datamodule)
199     trainer.test(model=task, datamodule=datamodule)
200 
201 
202 if __name__ == "__main__":
203     # Taken from https://github.com/pangeo-data/cog-best-practices
204     _rasterio_best_practices = {
205         "GDAL_DISABLE_READDIR_ON_OPEN": "EMPTY_DIR",
206         "AWS_NO_SIGN_REQUEST": "YES",
207         "GDAL_MAX_RAW_BLOCK_CACHE_SIZE": "200000000",
208         "GDAL_SWATH_SIZE": "200000000",
209         "VSI_CURL_CACHE_SIZE": "200000000",
210     }
211     os.environ.update(_rasterio_best_practices)
212 
213     conf = set_up_omegaconf()
214 
215     # Set random seed for reproducibility
216     # https://pytorch-lightning.readthedocs.io/en/latest/api/pytorch_lightning.utilities.seed.html#pytorch_lightning.utilities.seed.seed_everything
217     pl.seed_everything(conf.program.seed)
218 
219     # Main training procedure
220     main(conf)
221
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/train.py b/train.py
--- a/train.py
+++ b/train.py
@@ -39,28 +39,19 @@
 TASK_TO_MODULES_MAPPING: Dict[
     str, Tuple[Type[pl.LightningModule], Type[pl.LightningDataModule]]
 ] = {
-    "bigearthnet_all": (MultiLabelClassificationTask, BigEarthNetDataModule),
-    "bigearthnet_s1": (MultiLabelClassificationTask, BigEarthNetDataModule),
-    "bigearthnet_s2": (MultiLabelClassificationTask, BigEarthNetDataModule),
+    "bigearthnet": (MultiLabelClassificationTask, BigEarthNetDataModule),
     "byol": (BYOLTask, ChesapeakeCVPRDataModule),
-    "chesapeake_cvpr_5": (SemanticSegmentationTask, ChesapeakeCVPRDataModule),
-    "chesapeake_cvpr_7": (SemanticSegmentationTask, ChesapeakeCVPRDataModule),
-    "chesapeake_cvpr_prior": (SemanticSegmentationTask, ChesapeakeCVPRDataModule),
+    "chesapeake_cvpr": (SemanticSegmentationTask, ChesapeakeCVPRDataModule),
     "cowc_counting": (RegressionTask, COWCCountingDataModule),
     "cyclone": (RegressionTask, CycloneDataModule),
     "eurosat": (ClassificationTask, EuroSATDataModule),
     "etci2021": (SemanticSegmentationTask, ETCI2021DataModule),
     "landcoverai": (SemanticSegmentationTask, LandCoverAIDataModule),
     "naipchesapeake": (SemanticSegmentationTask, NAIPChesapeakeDataModule),
-    "oscd_all": (SemanticSegmentationTask, OSCDDataModule),
-    "oscd_rgb": (SemanticSegmentationTask, OSCDDataModule),
+    "oscd": (SemanticSegmentationTask, OSCDDataModule),
     "resisc45": (ClassificationTask, RESISC45DataModule),
-    "sen12ms_all": (SemanticSegmentationTask, SEN12MSDataModule),
-    "sen12ms_s1": (SemanticSegmentationTask, SEN12MSDataModule),
-    "sen12ms_s2_all": (SemanticSegmentationTask, SEN12MSDataModule),
-    "sen12ms_s2_reduced": (SemanticSegmentationTask, SEN12MSDataModule),
-    "so2sat_supervised": (ClassificationTask, So2SatDataModule),
-    "so2sat_unsupervised": (ClassificationTask, So2SatDataModule),
+    "sen12ms": (SemanticSegmentationTask, SEN12MSDataModule),
+    "so2sat": (ClassificationTask, So2SatDataModule),
     "ucmerced": (ClassificationTask, UCMercedDataModule),
 }
 
@@ -104,7 +95,7 @@
     # These OmegaConf structured configs enforce a schema at runtime, see:
     # https://omegaconf.readthedocs.io/en/2.0_branch/structured_config.html#merging-with-other-configs
     task_name = conf.experiment.task
-    task_config_fn = os.path.join("conf", "task_defaults", f"{task_name}.yaml")
+    task_config_fn = os.path.join("conf", f"{task_name}.yaml")
     if task_name == "test":
         task_conf = OmegaConf.create()
     elif os.path.exists(task_config_fn):

{"golden_diff": "diff --git a/train.py b/train.py\n--- a/train.py\n+++ b/train.py\n@@ -39,28 +39,19 @@\n TASK_TO_MODULES_MAPPING: Dict[\n     str, Tuple[Type[pl.LightningModule], Type[pl.LightningDataModule]]\n ] = {\n-    \"bigearthnet_all\": (MultiLabelClassificationTask, BigEarthNetDataModule),\n-    \"bigearthnet_s1\": (MultiLabelClassificationTask, BigEarthNetDataModule),\n-    \"bigearthnet_s2\": (MultiLabelClassificationTask, BigEarthNetDataModule),\n+    \"bigearthnet\": (MultiLabelClassificationTask, BigEarthNetDataModule),\n     \"byol\": (BYOLTask, ChesapeakeCVPRDataModule),\n-    \"chesapeake_cvpr_5\": (SemanticSegmentationTask, ChesapeakeCVPRDataModule),\n-    \"chesapeake_cvpr_7\": (SemanticSegmentationTask, ChesapeakeCVPRDataModule),\n-    \"chesapeake_cvpr_prior\": (SemanticSegmentationTask, ChesapeakeCVPRDataModule),\n+    \"chesapeake_cvpr\": (SemanticSegmentationTask, ChesapeakeCVPRDataModule),\n     \"cowc_counting\": (RegressionTask, COWCCountingDataModule),\n     \"cyclone\": (RegressionTask, CycloneDataModule),\n     \"eurosat\": (ClassificationTask, EuroSATDataModule),\n     \"etci2021\": (SemanticSegmentationTask, ETCI2021DataModule),\n     \"landcoverai\": (SemanticSegmentationTask, LandCoverAIDataModule),\n     \"naipchesapeake\": (SemanticSegmentationTask, NAIPChesapeakeDataModule),\n-    \"oscd_all\": (SemanticSegmentationTask, OSCDDataModule),\n-    \"oscd_rgb\": (SemanticSegmentationTask, OSCDDataModule),\n+    \"oscd\": (SemanticSegmentationTask, OSCDDataModule),\n     \"resisc45\": (ClassificationTask, RESISC45DataModule),\n-    \"sen12ms_all\": (SemanticSegmentationTask, SEN12MSDataModule),\n-    \"sen12ms_s1\": (SemanticSegmentationTask, SEN12MSDataModule),\n-    \"sen12ms_s2_all\": (SemanticSegmentationTask, SEN12MSDataModule),\n-    \"sen12ms_s2_reduced\": (SemanticSegmentationTask, SEN12MSDataModule),\n-    \"so2sat_supervised\": (ClassificationTask, So2SatDataModule),\n-    \"so2sat_unsupervised\": (ClassificationTask, So2SatDataModule),\n+    \"sen12ms\": (SemanticSegmentationTask, SEN12MSDataModule),\n+    \"so2sat\": (ClassificationTask, So2SatDataModule),\n     \"ucmerced\": (ClassificationTask, UCMercedDataModule),\n }\n \n@@ -104,7 +95,7 @@\n     # These OmegaConf structured configs enforce a schema at runtime, see:\n     # https://omegaconf.readthedocs.io/en/2.0_branch/structured_config.html#merging-with-other-configs\n     task_name = conf.experiment.task\n-    task_config_fn = os.path.join(\"conf\", \"task_defaults\", f\"{task_name}.yaml\")\n+    task_config_fn = os.path.join(\"conf\", f\"{task_name}.yaml\")\n     if task_name == \"test\":\n         task_conf = OmegaConf.create()\n     elif os.path.exists(task_config_fn):\n", "issue": "Re-think how configs are handled in train.py\nCurrently configuration to `train.py` is handled with [OmegaConf](https://omegaconf.readthedocs.io/en/2.1_branch/usage.html). This made more sense when the tasks (and accompanying trainer code) were fragmented, as we could easily define per-task configuration. Now that the trainer code that we would like to include in base TorchGeo are being generalized into things like `ClassificationTask` and `SemanticSegmentationTask` _and_ it is clear that more complicated training configurations won't be supported by torchgeo proper, it might make sense to pull out the OmegaConf part, and go with a more simple `argparse` based approach. Bonus: this would also allow us to get rid of a dependency. I'm not sure how exactly the argparse approach would work in all cases but it is worth more thought!\r\n\r\nLightning has a few pieces of docs that can help with this:\r\n- https://pytorch-lightning.readthedocs.io/en/stable/common/trainer.html#trainer-in-python-scripts\r\n- https://pytorch-lightning.readthedocs.io/en/stable/common/hyperparameters.html\r\n- https://pytorch-lightning.readthedocs.io/en/latest/common/lightning_cli.html\r\n\r\nWhatever we settle on here should definitely still allow passing arguments via a YAML config file. This allows reproducible benchmark experiment configurations to be saved in source control.\n", "before_files": [{"content": "#!/usr/bin/env python3\n\n# Copyright (c) Microsoft Corporation. All rights reserved.\n# Licensed under the MIT License.\n\n\"\"\"torchgeo model training script.\"\"\"\n\nimport os\nfrom typing import Any, Dict, Tuple, Type, cast\n\nimport pytorch_lightning as pl\nfrom omegaconf import DictConfig, OmegaConf\nfrom pytorch_lightning import loggers as pl_loggers\nfrom pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint\n\nfrom torchgeo.datamodules import (\n    BigEarthNetDataModule,\n    ChesapeakeCVPRDataModule,\n    COWCCountingDataModule,\n    CycloneDataModule,\n    ETCI2021DataModule,\n    EuroSATDataModule,\n    LandCoverAIDataModule,\n    NAIPChesapeakeDataModule,\n    OSCDDataModule,\n    RESISC45DataModule,\n    SEN12MSDataModule,\n    So2SatDataModule,\n    UCMercedDataModule,\n)\nfrom torchgeo.trainers import (\n    BYOLTask,\n    ClassificationTask,\n    MultiLabelClassificationTask,\n    RegressionTask,\n    SemanticSegmentationTask,\n)\n\nTASK_TO_MODULES_MAPPING: Dict[\n    str, Tuple[Type[pl.LightningModule], Type[pl.LightningDataModule]]\n] = {\n    \"bigearthnet_all\": (MultiLabelClassificationTask, BigEarthNetDataModule),\n    \"bigearthnet_s1\": (MultiLabelClassificationTask, BigEarthNetDataModule),\n    \"bigearthnet_s2\": (MultiLabelClassificationTask, BigEarthNetDataModule),\n    \"byol\": (BYOLTask, ChesapeakeCVPRDataModule),\n    \"chesapeake_cvpr_5\": (SemanticSegmentationTask, ChesapeakeCVPRDataModule),\n    \"chesapeake_cvpr_7\": (SemanticSegmentationTask, ChesapeakeCVPRDataModule),\n    \"chesapeake_cvpr_prior\": (SemanticSegmentationTask, ChesapeakeCVPRDataModule),\n    \"cowc_counting\": (RegressionTask, COWCCountingDataModule),\n    \"cyclone\": (RegressionTask, CycloneDataModule),\n    \"eurosat\": (ClassificationTask, EuroSATDataModule),\n    \"etci2021\": (SemanticSegmentationTask, ETCI2021DataModule),\n    \"landcoverai\": (SemanticSegmentationTask, LandCoverAIDataModule),\n    \"naipchesapeake\": (SemanticSegmentationTask, NAIPChesapeakeDataModule),\n    \"oscd_all\": (SemanticSegmentationTask, OSCDDataModule),\n    \"oscd_rgb\": (SemanticSegmentationTask, OSCDDataModule),\n    \"resisc45\": (ClassificationTask, RESISC45DataModule),\n    \"sen12ms_all\": (SemanticSegmentationTask, SEN12MSDataModule),\n    \"sen12ms_s1\": (SemanticSegmentationTask, SEN12MSDataModule),\n    \"sen12ms_s2_all\": (SemanticSegmentationTask, SEN12MSDataModule),\n    \"sen12ms_s2_reduced\": (SemanticSegmentationTask, SEN12MSDataModule),\n    \"so2sat_supervised\": (ClassificationTask, So2SatDataModule),\n    \"so2sat_unsupervised\": (ClassificationTask, So2SatDataModule),\n    \"ucmerced\": (ClassificationTask, UCMercedDataModule),\n}\n\n\ndef set_up_omegaconf() -> DictConfig:\n    \"\"\"Loads program arguments from either YAML config files or command line arguments.\n\n    This method loads defaults/a schema from \"conf/defaults.yaml\" as well as potential\n    arguments from the command line. If one of the command line arguments is\n    \"config_file\", then we additionally read arguments from that YAML file. One of the\n    config file based arguments or command line arguments must specify task.name. The\n    task.name value is used to grab a task specific defaults from its respective\n    trainer. The final configuration is given as merge(task_defaults, defaults,\n    config file, command line). The merge() works from the first argument to the last,\n    replacing existing values with newer values. Additionally, if any values are\n    merged into task_defaults without matching types, then there will be a runtime\n    error.\n\n    Returns:\n        an OmegaConf DictConfig containing all the validated program arguments\n\n    Raises:\n        FileNotFoundError: when ``config_file`` does not exist\n        ValueError: when ``task.name`` is not a valid task\n    \"\"\"\n    conf = OmegaConf.load(\"conf/defaults.yaml\")\n    command_line_conf = OmegaConf.from_cli()\n\n    if \"config_file\" in command_line_conf:\n        config_fn = command_line_conf.config_file\n        if not os.path.isfile(config_fn):\n            raise FileNotFoundError(f\"config_file={config_fn} is not a valid file\")\n\n        user_conf = OmegaConf.load(config_fn)\n        conf = OmegaConf.merge(conf, user_conf)\n\n    conf = OmegaConf.merge(  # Merge in any arguments passed via the command line\n        conf, command_line_conf\n    )\n\n    # These OmegaConf structured configs enforce a schema at runtime, see:\n    # https://omegaconf.readthedocs.io/en/2.0_branch/structured_config.html#merging-with-other-configs\n    task_name = conf.experiment.task\n    task_config_fn = os.path.join(\"conf\", \"task_defaults\", f\"{task_name}.yaml\")\n    if task_name == \"test\":\n        task_conf = OmegaConf.create()\n    elif os.path.exists(task_config_fn):\n        task_conf = cast(DictConfig, OmegaConf.load(task_config_fn))\n    else:\n        raise ValueError(\n            f\"experiment.task={task_name} is not recognized as a valid task\"\n        )\n\n    conf = OmegaConf.merge(task_conf, conf)\n    conf = cast(DictConfig, conf)  # convince mypy that everything is alright\n\n    return conf\n\n\ndef main(conf: DictConfig) -> None:\n    \"\"\"Main training loop.\"\"\"\n    ######################################\n    # Setup output directory\n    ######################################\n\n    experiment_name = conf.experiment.name\n    task_name = conf.experiment.task\n    if os.path.isfile(conf.program.output_dir):\n        raise NotADirectoryError(\"`program.output_dir` must be a directory\")\n    os.makedirs(conf.program.output_dir, exist_ok=True)\n\n    experiment_dir = os.path.join(conf.program.output_dir, experiment_name)\n    os.makedirs(experiment_dir, exist_ok=True)\n\n    if len(os.listdir(experiment_dir)) > 0:\n        if conf.program.overwrite:\n            print(\n                f\"WARNING! The experiment directory, {experiment_dir}, already exists, \"\n                + \"we might overwrite data in it!\"\n            )\n        else:\n            raise FileExistsError(\n                f\"The experiment directory, {experiment_dir}, already exists and isn't \"\n                + \"empty. We don't want to overwrite any existing results, exiting...\"\n            )\n\n    with open(os.path.join(experiment_dir, \"experiment_config.yaml\"), \"w\") as f:\n        OmegaConf.save(config=conf, f=f)\n\n    ######################################\n    # Choose task to run based on arguments or configuration\n    ######################################\n    # Convert the DictConfig into a dictionary so that we can pass as kwargs.\n    task_args = cast(Dict[str, Any], OmegaConf.to_object(conf.experiment.module))\n    datamodule_args = cast(\n        Dict[str, Any], OmegaConf.to_object(conf.experiment.datamodule)\n    )\n\n    datamodule: pl.LightningDataModule\n    task: pl.LightningModule\n    if task_name in TASK_TO_MODULES_MAPPING:\n        task_class, datamodule_class = TASK_TO_MODULES_MAPPING[task_name]\n        task = task_class(**task_args)\n        datamodule = datamodule_class(**datamodule_args)\n    else:\n        raise ValueError(\n            f\"experiment.task={task_name} is not recognized as a valid task\"\n        )\n\n    ######################################\n    # Setup trainer\n    ######################################\n    tb_logger = pl_loggers.TensorBoardLogger(conf.program.log_dir, name=experiment_name)\n\n    checkpoint_callback = ModelCheckpoint(\n        monitor=\"val_loss\", dirpath=experiment_dir, save_top_k=1, save_last=True\n    )\n    early_stopping_callback = EarlyStopping(\n        monitor=\"val_loss\", min_delta=0.00, patience=18\n    )\n\n    trainer_args = cast(Dict[str, Any], OmegaConf.to_object(conf.trainer))\n\n    trainer_args[\"callbacks\"] = [checkpoint_callback, early_stopping_callback]\n    trainer_args[\"logger\"] = tb_logger\n    trainer_args[\"default_root_dir\"] = experiment_dir\n    trainer = pl.Trainer(**trainer_args)\n\n    if trainer_args.get(\"auto_lr_find\"):\n        trainer.tune(model=task, datamodule=datamodule)\n\n    ######################################\n    # Run experiment\n    ######################################\n    trainer.fit(model=task, datamodule=datamodule)\n    trainer.test(model=task, datamodule=datamodule)\n\n\nif __name__ == \"__main__\":\n    # Taken from https://github.com/pangeo-data/cog-best-practices\n    _rasterio_best_practices = {\n        \"GDAL_DISABLE_READDIR_ON_OPEN\": \"EMPTY_DIR\",\n        \"AWS_NO_SIGN_REQUEST\": \"YES\",\n        \"GDAL_MAX_RAW_BLOCK_CACHE_SIZE\": \"200000000\",\n        \"GDAL_SWATH_SIZE\": \"200000000\",\n        \"VSI_CURL_CACHE_SIZE\": \"200000000\",\n    }\n    os.environ.update(_rasterio_best_practices)\n\n    conf = set_up_omegaconf()\n\n    # Set random seed for reproducibility\n    # https://pytorch-lightning.readthedocs.io/en/latest/api/pytorch_lightning.utilities.seed.html#pytorch_lightning.utilities.seed.seed_everything\n    pl.seed_everything(conf.program.seed)\n\n    # Main training procedure\n    main(conf)\n", "path": "train.py"}], "after_files": [{"content": "#!/usr/bin/env python3\n\n# Copyright (c) Microsoft Corporation. All rights reserved.\n# Licensed under the MIT License.\n\n\"\"\"torchgeo model training script.\"\"\"\n\nimport os\nfrom typing import Any, Dict, Tuple, Type, cast\n\nimport pytorch_lightning as pl\nfrom omegaconf import DictConfig, OmegaConf\nfrom pytorch_lightning import loggers as pl_loggers\nfrom pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint\n\nfrom torchgeo.datamodules import (\n    BigEarthNetDataModule,\n    ChesapeakeCVPRDataModule,\n    COWCCountingDataModule,\n    CycloneDataModule,\n    ETCI2021DataModule,\n    EuroSATDataModule,\n    LandCoverAIDataModule,\n    NAIPChesapeakeDataModule,\n    OSCDDataModule,\n    RESISC45DataModule,\n    SEN12MSDataModule,\n    So2SatDataModule,\n    UCMercedDataModule,\n)\nfrom torchgeo.trainers import (\n    BYOLTask,\n    ClassificationTask,\n    MultiLabelClassificationTask,\n    RegressionTask,\n    SemanticSegmentationTask,\n)\n\nTASK_TO_MODULES_MAPPING: Dict[\n    str, Tuple[Type[pl.LightningModule], Type[pl.LightningDataModule]]\n] = {\n    \"bigearthnet\": (MultiLabelClassificationTask, BigEarthNetDataModule),\n    \"byol\": (BYOLTask, ChesapeakeCVPRDataModule),\n    \"chesapeake_cvpr\": (SemanticSegmentationTask, ChesapeakeCVPRDataModule),\n    \"cowc_counting\": (RegressionTask, COWCCountingDataModule),\n    \"cyclone\": (RegressionTask, CycloneDataModule),\n    \"eurosat\": (ClassificationTask, EuroSATDataModule),\n    \"etci2021\": (SemanticSegmentationTask, ETCI2021DataModule),\n    \"landcoverai\": (SemanticSegmentationTask, LandCoverAIDataModule),\n    \"naipchesapeake\": (SemanticSegmentationTask, NAIPChesapeakeDataModule),\n    \"oscd\": (SemanticSegmentationTask, OSCDDataModule),\n    \"resisc45\": (ClassificationTask, RESISC45DataModule),\n    \"sen12ms\": (SemanticSegmentationTask, SEN12MSDataModule),\n    \"so2sat\": (ClassificationTask, So2SatDataModule),\n    \"ucmerced\": (ClassificationTask, UCMercedDataModule),\n}\n\n\ndef set_up_omegaconf() -> DictConfig:\n    \"\"\"Loads program arguments from either YAML config files or command line arguments.\n\n    This method loads defaults/a schema from \"conf/defaults.yaml\" as well as potential\n    arguments from the command line. If one of the command line arguments is\n    \"config_file\", then we additionally read arguments from that YAML file. One of the\n    config file based arguments or command line arguments must specify task.name. The\n    task.name value is used to grab a task specific defaults from its respective\n    trainer. The final configuration is given as merge(task_defaults, defaults,\n    config file, command line). The merge() works from the first argument to the last,\n    replacing existing values with newer values. Additionally, if any values are\n    merged into task_defaults without matching types, then there will be a runtime\n    error.\n\n    Returns:\n        an OmegaConf DictConfig containing all the validated program arguments\n\n    Raises:\n        FileNotFoundError: when ``config_file`` does not exist\n        ValueError: when ``task.name`` is not a valid task\n    \"\"\"\n    conf = OmegaConf.load(\"conf/defaults.yaml\")\n    command_line_conf = OmegaConf.from_cli()\n\n    if \"config_file\" in command_line_conf:\n        config_fn = command_line_conf.config_file\n        if not os.path.isfile(config_fn):\n            raise FileNotFoundError(f\"config_file={config_fn} is not a valid file\")\n\n        user_conf = OmegaConf.load(config_fn)\n        conf = OmegaConf.merge(conf, user_conf)\n\n    conf = OmegaConf.merge(  # Merge in any arguments passed via the command line\n        conf, command_line_conf\n    )\n\n    # These OmegaConf structured configs enforce a schema at runtime, see:\n    # https://omegaconf.readthedocs.io/en/2.0_branch/structured_config.html#merging-with-other-configs\n    task_name = conf.experiment.task\n    task_config_fn = os.path.join(\"conf\", f\"{task_name}.yaml\")\n    if task_name == \"test\":\n        task_conf = OmegaConf.create()\n    elif os.path.exists(task_config_fn):\n        task_conf = cast(DictConfig, OmegaConf.load(task_config_fn))\n    else:\n        raise ValueError(\n            f\"experiment.task={task_name} is not recognized as a valid task\"\n        )\n\n    conf = OmegaConf.merge(task_conf, conf)\n    conf = cast(DictConfig, conf)  # convince mypy that everything is alright\n\n    return conf\n\n\ndef main(conf: DictConfig) -> None:\n    \"\"\"Main training loop.\"\"\"\n    ######################################\n    # Setup output directory\n    ######################################\n\n    experiment_name = conf.experiment.name\n    task_name = conf.experiment.task\n    if os.path.isfile(conf.program.output_dir):\n        raise NotADirectoryError(\"`program.output_dir` must be a directory\")\n    os.makedirs(conf.program.output_dir, exist_ok=True)\n\n    experiment_dir = os.path.join(conf.program.output_dir, experiment_name)\n    os.makedirs(experiment_dir, exist_ok=True)\n\n    if len(os.listdir(experiment_dir)) > 0:\n        if conf.program.overwrite:\n            print(\n                f\"WARNING! The experiment directory, {experiment_dir}, already exists, \"\n                + \"we might overwrite data in it!\"\n            )\n        else:\n            raise FileExistsError(\n                f\"The experiment directory, {experiment_dir}, already exists and isn't \"\n                + \"empty. We don't want to overwrite any existing results, exiting...\"\n            )\n\n    with open(os.path.join(experiment_dir, \"experiment_config.yaml\"), \"w\") as f:\n        OmegaConf.save(config=conf, f=f)\n\n    ######################################\n    # Choose task to run based on arguments or configuration\n    ######################################\n    # Convert the DictConfig into a dictionary so that we can pass as kwargs.\n    task_args = cast(Dict[str, Any], OmegaConf.to_object(conf.experiment.module))\n    datamodule_args = cast(\n        Dict[str, Any], OmegaConf.to_object(conf.experiment.datamodule)\n    )\n\n    datamodule: pl.LightningDataModule\n    task: pl.LightningModule\n    if task_name in TASK_TO_MODULES_MAPPING:\n        task_class, datamodule_class = TASK_TO_MODULES_MAPPING[task_name]\n        task = task_class(**task_args)\n        datamodule = datamodule_class(**datamodule_args)\n    else:\n        raise ValueError(\n            f\"experiment.task={task_name} is not recognized as a valid task\"\n        )\n\n    ######################################\n    # Setup trainer\n    ######################################\n    tb_logger = pl_loggers.TensorBoardLogger(conf.program.log_dir, name=experiment_name)\n\n    checkpoint_callback = ModelCheckpoint(\n        monitor=\"val_loss\", dirpath=experiment_dir, save_top_k=1, save_last=True\n    )\n    early_stopping_callback = EarlyStopping(\n        monitor=\"val_loss\", min_delta=0.00, patience=18\n    )\n\n    trainer_args = cast(Dict[str, Any], OmegaConf.to_object(conf.trainer))\n\n    trainer_args[\"callbacks\"] = [checkpoint_callback, early_stopping_callback]\n    trainer_args[\"logger\"] = tb_logger\n    trainer_args[\"default_root_dir\"] = experiment_dir\n    trainer = pl.Trainer(**trainer_args)\n\n    if trainer_args.get(\"auto_lr_find\"):\n        trainer.tune(model=task, datamodule=datamodule)\n\n    ######################################\n    # Run experiment\n    ######################################\n    trainer.fit(model=task, datamodule=datamodule)\n    trainer.test(model=task, datamodule=datamodule)\n\n\nif __name__ == \"__main__\":\n    # Taken from https://github.com/pangeo-data/cog-best-practices\n    _rasterio_best_practices = {\n        \"GDAL_DISABLE_READDIR_ON_OPEN\": \"EMPTY_DIR\",\n        \"AWS_NO_SIGN_REQUEST\": \"YES\",\n        \"GDAL_MAX_RAW_BLOCK_CACHE_SIZE\": \"200000000\",\n        \"GDAL_SWATH_SIZE\": \"200000000\",\n        \"VSI_CURL_CACHE_SIZE\": \"200000000\",\n    }\n    os.environ.update(_rasterio_best_practices)\n\n    conf = set_up_omegaconf()\n\n    # Set random seed for reproducibility\n    # https://pytorch-lightning.readthedocs.io/en/latest/api/pytorch_lightning.utilities.seed.html#pytorch_lightning.utilities.seed.seed_everything\n    pl.seed_everything(conf.program.seed)\n\n    # Main training procedure\n    main(conf)\n", "path": "train.py"}]}