We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Add CSVLogger to train.py
### Summary

Our current `train.py` script only logs to TensorBoard. We should add a CSV logger.

### Rationale

In order to determine the optimal hyperparams in a hyperparam sweep, we need an easy way of finding the final val accuracy of each run. Currently, the only way to get this is to convert the TensorBoard logger output to a CSV. If we log directly to CSV this becomes much easier, especially for those who don't use TensorBoard.

### Implementation

_No response_

### Alternatives

_No response_

### Additional information

_No response_

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

{"golden_diff": "diff --git a/train.py b/train.py\n--- a/train.py\n+++ b/train.py\n@@ -170,6 +170,7 @@\n     # Setup trainer\n     ######################################\n     tb_logger = pl_loggers.TensorBoardLogger(conf.program.log_dir, name=experiment_name)\n+    csv_logger = pl_loggers.CSVLogger(conf.program.log_dir, name=experiment_name)\n \n     if isinstance(task, ObjectDetectionTask):\n         monitor_metric = \"val_map\"\n@@ -192,7 +193,7 @@\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[\"logger\"] = [tb_logger, csv_logger]\n     trainer_args[\"default_root_dir\"] = experiment_dir\n     trainer = pl.Trainer(**trainer_args)\n", "issue": "Add CSVLogger to train.py\n### Summary\n\nOur current `train.py` script only logs to TensorBoard. We should add a CSV logger.\n\n### Rationale\n\nIn order to determine the optimal hyperparams in a hyperparam sweep, we need an easy way of finding the final val accuracy of each run. Currently, the only way to get this is to convert the TensorBoard logger output to a CSV. If we log directly to CSV this becomes much easier, especially for those who don't use TensorBoard.\n\n### Implementation\n\n_No response_\n\n### Alternatives\n\n_No response_\n\n### Additional information\n\n_No response_\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    ETCI2021DataModule,\n    EuroSATDataModule,\n    InriaAerialImageLabelingDataModule,\n    LandCoverAIDataModule,\n    NAIPChesapeakeDataModule,\n    NASAMarineDebrisDataModule,\n    RESISC45DataModule,\n    SEN12MSDataModule,\n    So2SatDataModule,\n    SpaceNet1DataModule,\n    TropicalCycloneDataModule,\n    UCMercedDataModule,\n)\nfrom torchgeo.trainers import (\n    BYOLTask,\n    ClassificationTask,\n    MultiLabelClassificationTask,\n    ObjectDetectionTask,\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, TropicalCycloneDataModule),\n    \"eurosat\": (ClassificationTask, EuroSATDataModule),\n    \"etci2021\": (SemanticSegmentationTask, ETCI2021DataModule),\n    \"inria\": (SemanticSegmentationTask, InriaAerialImageLabelingDataModule),\n    \"landcoverai\": (SemanticSegmentationTask, LandCoverAIDataModule),\n    \"naipchesapeake\": (SemanticSegmentationTask, NAIPChesapeakeDataModule),\n    \"nasa_marine_debris\": (ObjectDetectionTask, NASAMarineDebrisDataModule),\n    \"resisc45\": (ClassificationTask, RESISC45DataModule),\n    \"sen12ms\": (SemanticSegmentationTask, SEN12MSDataModule),\n    \"so2sat\": (ClassificationTask, So2SatDataModule),\n    \"spacenet1\": (SemanticSegmentationTask, SpaceNet1DataModule),\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    if isinstance(task, ObjectDetectionTask):\n        monitor_metric = \"val_map\"\n        mode = \"max\"\n    else:\n        monitor_metric = \"val_loss\"\n        mode = \"min\"\n\n    checkpoint_callback = ModelCheckpoint(\n        monitor=monitor_metric,\n        filename=\"checkpoint-epoch{epoch:02d}-val_loss{val_loss:.2f}\",\n        dirpath=experiment_dir,\n        save_top_k=1,\n        save_last=True,\n    )\n    early_stopping_callback = EarlyStopping(\n        monitor=monitor_metric, min_delta=0.00, patience=18, mode=mode\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(ckpt_path=\"best\", 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    ETCI2021DataModule,\n    EuroSATDataModule,\n    InriaAerialImageLabelingDataModule,\n    LandCoverAIDataModule,\n    NAIPChesapeakeDataModule,\n    NASAMarineDebrisDataModule,\n    RESISC45DataModule,\n    SEN12MSDataModule,\n    So2SatDataModule,\n    SpaceNet1DataModule,\n    TropicalCycloneDataModule,\n    UCMercedDataModule,\n)\nfrom torchgeo.trainers import (\n    BYOLTask,\n    ClassificationTask,\n    MultiLabelClassificationTask,\n    ObjectDetectionTask,\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, TropicalCycloneDataModule),\n    \"eurosat\": (ClassificationTask, EuroSATDataModule),\n    \"etci2021\": (SemanticSegmentationTask, ETCI2021DataModule),\n    \"inria\": (SemanticSegmentationTask, InriaAerialImageLabelingDataModule),\n    \"landcoverai\": (SemanticSegmentationTask, LandCoverAIDataModule),\n    \"naipchesapeake\": (SemanticSegmentationTask, NAIPChesapeakeDataModule),\n    \"nasa_marine_debris\": (ObjectDetectionTask, NASAMarineDebrisDataModule),\n    \"resisc45\": (ClassificationTask, RESISC45DataModule),\n    \"sen12ms\": (SemanticSegmentationTask, SEN12MSDataModule),\n    \"so2sat\": (ClassificationTask, So2SatDataModule),\n    \"spacenet1\": (SemanticSegmentationTask, SpaceNet1DataModule),\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    csv_logger = pl_loggers.CSVLogger(conf.program.log_dir, name=experiment_name)\n\n    if isinstance(task, ObjectDetectionTask):\n        monitor_metric = \"val_map\"\n        mode = \"max\"\n    else:\n        monitor_metric = \"val_loss\"\n        mode = \"min\"\n\n    checkpoint_callback = ModelCheckpoint(\n        monitor=monitor_metric,\n        filename=\"checkpoint-epoch{epoch:02d}-val_loss{val_loss:.2f}\",\n        dirpath=experiment_dir,\n        save_top_k=1,\n        save_last=True,\n    )\n    early_stopping_callback = EarlyStopping(\n        monitor=monitor_metric, min_delta=0.00, patience=18, mode=mode\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, csv_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(ckpt_path=\"best\", 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"}]}