We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
CPU Batch jobs cannot pull container
After making a large number of dependency upgrades suggested by Dependabot, I started getting the following error when running any job that runs in the CPU compute environment on Batch. The GPU ones work fine. 

See https://aws.amazon.com/premiumsupport/knowledge-center/batch-job-failure-disk-space/

```
CannotPullContainerError: failed to register layer: ApplyLayer exit status 1 stdout: stderr: write /root/.cache/pip/http/d/c/1/1/c/dc11c115dbb602e6636493addd92f61ac957da65f9166a69f6820fbc: no space left on device
```

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `rastervision_pytorch_learner/rastervision/pytorch_learner/semantic_segmentation_learner.py`
Content:
```
1 import warnings
2 warnings.filterwarnings('ignore')  # noqa
3 
4 from typing import Union, Iterable, Optional
5 
6 import logging
7 
8 import numpy as np
9 import matplotlib
10 from matplotlib import pyplot as plt
11 import matplotlib.patches as mpatches
12 matplotlib.use('Agg')  # noqa
13 import albumentations as A
14 
15 import torch
16 from torch import nn
17 from torch.nn import functional as F
18 from torchvision import models
19 
20 from rastervision.pipeline.config import ConfigError
21 from rastervision.pytorch_learner.learner import Learner
22 from rastervision.pytorch_learner.utils import (
23     compute_conf_mat_metrics, compute_conf_mat, color_to_triple, SplitTensor,
24     Parallel, AddTensors)
25 from rastervision.pipeline.file_system import make_dir
26 
27 log = logging.getLogger(__name__)
28 
29 
30 class SemanticSegmentationLearner(Learner):
31     def build_model(self) -> nn.Module:
32         # TODO support FCN option
33         pretrained = self.cfg.model.pretrained
34         out_classes = len(self.cfg.data.class_names)
35         if self.cfg.solver.ignore_last_class:
36             out_classes -= 1
37         model = models.segmentation.segmentation._segm_resnet(
38             'deeplabv3',
39             self.cfg.model.get_backbone_str(),
40             out_classes,
41             False,
42             pretrained_backbone=pretrained)
43 
44         input_channels = self.cfg.data.img_channels
45         old_conv = model.backbone.conv1
46 
47         if input_channels == old_conv.in_channels:
48             return model
49 
50         # these parameters will be the same for the new conv layer
51         old_conv_args = {
52             'out_channels': old_conv.out_channels,
53             'kernel_size': old_conv.kernel_size,
54             'stride': old_conv.stride,
55             'padding': old_conv.padding,
56             'dilation': old_conv.dilation,
57             'groups': old_conv.groups,
58             'bias': old_conv.bias
59         }
60 
61         if not pretrained:
62             # simply replace the first conv layer with one with the
63             # correct number of input channels
64             new_conv = nn.Conv2d(in_channels=input_channels, **old_conv_args)
65             model.backbone.conv1 = new_conv
66             return model
67 
68         if input_channels > old_conv.in_channels:
69             # insert a new conv layer parallel to the existing one
70             # and sum their outputs
71             new_conv_channels = input_channels - old_conv.in_channels
72             new_conv = nn.Conv2d(
73                 in_channels=new_conv_channels, **old_conv_args)
74             model.backbone.conv1 = nn.Sequential(
75                 # split input along channel dim
76                 SplitTensor((old_conv.in_channels, new_conv_channels), dim=1),
77                 # each split goes to its respective conv layer
78                 Parallel(old_conv, new_conv),
79                 # sum the parallel outputs
80                 AddTensors())
81         elif input_channels < old_conv.in_channels:
82             model.backbone.conv1 = nn.Conv2d(
83                 in_channels=input_channels, **old_conv_args)
84             model.backbone.conv1.weight.data[:, :input_channels] = \
85                 old_conv.weight.data[:, :input_channels]
86         else:
87             raise ConfigError(f'Something went wrong')
88 
89         return model
90 
91     def build_loss(self):
92         args = {}
93 
94         loss_weights = self.cfg.solver.class_loss_weights
95         if loss_weights is not None:
96             loss_weights = torch.tensor(loss_weights, device=self.device)
97             args.update({'weight': loss_weights})
98 
99         if self.cfg.solver.ignore_last_class:
100             num_classes = len(self.cfg.data.class_names)
101             args.update({'ignore_index': num_classes - 1})
102 
103         loss = nn.CrossEntropyLoss(**args)
104 
105         return loss
106 
107     def train_step(self, batch, batch_ind):
108         x, y = batch
109         out = self.post_forward(self.model(x))
110         return {'train_loss': self.loss(out, y)}
111 
112     def validate_step(self, batch, batch_ind):
113         x, y = batch
114         out = self.post_forward(self.model(x))
115         val_loss = self.loss(out, y)
116 
117         num_labels = len(self.cfg.data.class_names)
118         y = y.view(-1)
119         out = self.prob_to_pred(out).view(-1)
120         conf_mat = compute_conf_mat(out, y, num_labels)
121 
122         return {'val_loss': val_loss, 'conf_mat': conf_mat}
123 
124     def validate_end(self, outputs, num_samples):
125         conf_mat = sum([o['conf_mat'] for o in outputs])
126         val_loss = torch.stack([o['val_loss']
127                                 for o in outputs]).sum() / num_samples
128         conf_mat_metrics = compute_conf_mat_metrics(conf_mat,
129                                                     self.cfg.data.class_names)
130 
131         metrics = {'val_loss': val_loss.item()}
132         metrics.update(conf_mat_metrics)
133 
134         return metrics
135 
136     def post_forward(self, x):
137         if isinstance(x, dict):
138             return x['out']
139         return x
140 
141     def predict(self, x: torch.Tensor, raw_out: bool = False) -> torch.Tensor:
142         x = self.to_batch(x).float()
143         x = self.to_device(x, self.device)
144         with torch.no_grad():
145             out = self.model(x)
146             out = self.post_forward(out)
147             out = out.softmax(dim=1)
148             if not raw_out:
149                 out = self.prob_to_pred(out)
150         out = self.to_device(out, 'cpu')
151         return out
152 
153     def numpy_predict(self, x: np.ndarray,
154                       raw_out: bool = False) -> np.ndarray:
155         _, h, w, _ = x.shape
156         transform, _ = self.get_data_transforms()
157         x = self.normalize_input(x)
158         x = self.to_batch(x)
159         x = np.stack([transform(image=img)['image'] for img in x])
160         x = torch.from_numpy(x)
161         x = x.permute((0, 3, 1, 2))
162         out = self.predict(x, raw_out=True)
163         out = F.interpolate(
164             out, size=(h, w), mode='bilinear', align_corners=False)
165         out = self.prob_to_pred(out)
166         return self.output_to_numpy(out)
167 
168     def prob_to_pred(self, x):
169         return x.argmax(1)
170 
171     def plot_batch(self,
172                    x: torch.Tensor,
173                    y: Union[torch.Tensor, np.ndarray],
174                    output_path: str,
175                    z: Optional[torch.Tensor] = None,
176                    batch_limit: Optional[int] = None) -> None:
177         """Plot a whole batch in a grid using plot_xyz.
178 
179         Args:
180             x: batch of images
181             y: ground truth labels
182             output_path: local path where to save plot image
183             z: optional predicted labels
184             batch_limit: optional limit on (rendered) batch size
185         """
186         batch_sz, c, h, w = x.shape
187         batch_sz = min(batch_sz,
188                        batch_limit) if batch_limit is not None else batch_sz
189         if batch_sz == 0:
190             return
191 
192         channel_groups = self.cfg.data.channel_display_groups
193 
194         nrows = batch_sz
195         # one col for each group + 1 for labels + 1 for predictions
196         ncols = len(channel_groups) + 1
197         if z is not None:
198             ncols += 1
199 
200         fig, axes = plt.subplots(
201             nrows=nrows,
202             ncols=ncols,
203             squeeze=False,
204             constrained_layout=True,
205             figsize=(3 * ncols, 3 * nrows))
206 
207         assert axes.shape == (nrows, ncols)
208 
209         # (N, c, h, w) --> (N, h, w, c)
210         x = x.permute(0, 2, 3, 1)
211 
212         # apply transform, if given
213         if self.cfg.data.plot_options.transform is not None:
214             tf = A.from_dict(self.cfg.data.plot_options.transform)
215             imgs = [tf(image=img)['image'] for img in x.numpy()]
216             x = torch.from_numpy(np.stack(imgs))
217 
218         for i in range(batch_sz):
219             ax = (fig, axes[i])
220             if z is None:
221                 self.plot_xyz(ax, x[i], y[i])
222             else:
223                 self.plot_xyz(ax, x[i], y[i], z=z[i])
224 
225         make_dir(output_path, use_dirname=True)
226         plt.savefig(output_path, bbox_inches='tight')
227         plt.close()
228 
229     def plot_xyz(self,
230                  ax: Iterable,
231                  x: torch.Tensor,
232                  y: Union[torch.Tensor, np.ndarray],
233                  z: Optional[torch.Tensor] = None) -> None:
234 
235         channel_groups = self.cfg.data.channel_display_groups
236 
237         # make subplot titles
238         if not isinstance(channel_groups, dict):
239             channel_groups = {
240                 f'Channels: {[*chs]}': chs
241                 for chs in channel_groups
242             }
243 
244         fig, ax = ax
245         img_axes = ax[:len(channel_groups)]
246         label_ax = ax[len(channel_groups)]
247 
248         # plot input image(s)
249         for (title, chs), ch_ax in zip(channel_groups.items(), img_axes):
250             im = x[..., chs]
251             if len(chs) == 1:
252                 # repeat single channel 3 times
253                 im = im.expand(-1, -1, 3)
254             elif len(chs) == 2:
255                 # add a 3rd channel with all pixels set to 0.5
256                 h, w, _ = x.shape
257                 third_channel = torch.full((h, w, 1), fill_value=.5)
258                 im = torch.cat((im, third_channel), dim=-1)
259             elif len(chs) > 3:
260                 # only use the first 3 channels
261                 log.warn(f'Only plotting first 3 channels of channel-group '
262                          f'{title}: {chs}.')
263                 im = x[..., chs[:3]]
264             ch_ax.imshow(im)
265             ch_ax.set_title(title)
266             ch_ax.set_xticks([])
267             ch_ax.set_yticks([])
268 
269         class_colors = self.cfg.data.class_colors
270         colors = [color_to_triple(c) for c in class_colors]
271         colors = np.array(colors) / 255.
272         cmap = matplotlib.colors.ListedColormap(colors)
273 
274         # plot labels
275         label_ax.imshow(
276             y, vmin=0, vmax=len(colors), cmap=cmap, interpolation='none')
277         label_ax.set_title(f'Ground truth labels')
278         label_ax.set_xticks([])
279         label_ax.set_yticks([])
280 
281         # plot predictions
282         if z is not None:
283             pred_ax = ax[-1]
284             pred_ax.imshow(
285                 z, vmin=0, vmax=len(colors), cmap=cmap, interpolation='none')
286             pred_ax.set_title(f'Predicted labels')
287             pred_ax.set_xticks([])
288             pred_ax.set_yticks([])
289 
290         # add a legend to the rightmost subplot
291         class_names = self.cfg.data.class_names
292         legend_items = [
293             mpatches.Patch(facecolor=col, edgecolor='black', label=name)
294             for col, name in zip(colors, class_names)
295         ]
296         ax[-1].legend(
297             handles=legend_items,
298             loc='center right',
299             bbox_to_anchor=(1.8, 0.5))
300
```
--- 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/rastervision_pytorch_learner/rastervision/pytorch_learner/semantic_segmentation_learner.py b/rastervision_pytorch_learner/rastervision/pytorch_learner/semantic_segmentation_learner.py
--- a/rastervision_pytorch_learner/rastervision/pytorch_learner/semantic_segmentation_learner.py
+++ b/rastervision_pytorch_learner/rastervision/pytorch_learner/semantic_segmentation_learner.py
@@ -34,7 +34,7 @@
         out_classes = len(self.cfg.data.class_names)
         if self.cfg.solver.ignore_last_class:
             out_classes -= 1
-        model = models.segmentation.segmentation._segm_resnet(
+        model = models.segmentation.segmentation._segm_model(
             'deeplabv3',
             self.cfg.model.get_backbone_str(),
             out_classes,

{"golden_diff": "diff --git a/rastervision_pytorch_learner/rastervision/pytorch_learner/semantic_segmentation_learner.py b/rastervision_pytorch_learner/rastervision/pytorch_learner/semantic_segmentation_learner.py\n--- a/rastervision_pytorch_learner/rastervision/pytorch_learner/semantic_segmentation_learner.py\n+++ b/rastervision_pytorch_learner/rastervision/pytorch_learner/semantic_segmentation_learner.py\n@@ -34,7 +34,7 @@\n         out_classes = len(self.cfg.data.class_names)\n         if self.cfg.solver.ignore_last_class:\n             out_classes -= 1\n-        model = models.segmentation.segmentation._segm_resnet(\n+        model = models.segmentation.segmentation._segm_model(\n             'deeplabv3',\n             self.cfg.model.get_backbone_str(),\n             out_classes,\n", "issue": "CPU Batch jobs cannot pull container\nAfter making a large number of dependency upgrades suggested by Dependabot, I started getting the following error when running any job that runs in the CPU compute environment on Batch. The GPU ones work fine. \r\n\r\nSee https://aws.amazon.com/premiumsupport/knowledge-center/batch-job-failure-disk-space/\r\n\r\n```\r\nCannotPullContainerError: failed to register layer: ApplyLayer exit status 1 stdout: stderr: write /root/.cache/pip/http/d/c/1/1/c/dc11c115dbb602e6636493addd92f61ac957da65f9166a69f6820fbc: no space left on device\r\n```\n", "before_files": [{"content": "import warnings\nwarnings.filterwarnings('ignore')  # noqa\n\nfrom typing import Union, Iterable, Optional\n\nimport logging\n\nimport numpy as np\nimport matplotlib\nfrom matplotlib import pyplot as plt\nimport matplotlib.patches as mpatches\nmatplotlib.use('Agg')  # noqa\nimport albumentations as A\n\nimport torch\nfrom torch import nn\nfrom torch.nn import functional as F\nfrom torchvision import models\n\nfrom rastervision.pipeline.config import ConfigError\nfrom rastervision.pytorch_learner.learner import Learner\nfrom rastervision.pytorch_learner.utils import (\n    compute_conf_mat_metrics, compute_conf_mat, color_to_triple, SplitTensor,\n    Parallel, AddTensors)\nfrom rastervision.pipeline.file_system import make_dir\n\nlog = logging.getLogger(__name__)\n\n\nclass SemanticSegmentationLearner(Learner):\n    def build_model(self) -> nn.Module:\n        # TODO support FCN option\n        pretrained = self.cfg.model.pretrained\n        out_classes = len(self.cfg.data.class_names)\n        if self.cfg.solver.ignore_last_class:\n            out_classes -= 1\n        model = models.segmentation.segmentation._segm_resnet(\n            'deeplabv3',\n            self.cfg.model.get_backbone_str(),\n            out_classes,\n            False,\n            pretrained_backbone=pretrained)\n\n        input_channels = self.cfg.data.img_channels\n        old_conv = model.backbone.conv1\n\n        if input_channels == old_conv.in_channels:\n            return model\n\n        # these parameters will be the same for the new conv layer\n        old_conv_args = {\n            'out_channels': old_conv.out_channels,\n            'kernel_size': old_conv.kernel_size,\n            'stride': old_conv.stride,\n            'padding': old_conv.padding,\n            'dilation': old_conv.dilation,\n            'groups': old_conv.groups,\n            'bias': old_conv.bias\n        }\n\n        if not pretrained:\n            # simply replace the first conv layer with one with the\n            # correct number of input channels\n            new_conv = nn.Conv2d(in_channels=input_channels, **old_conv_args)\n            model.backbone.conv1 = new_conv\n            return model\n\n        if input_channels > old_conv.in_channels:\n            # insert a new conv layer parallel to the existing one\n            # and sum their outputs\n            new_conv_channels = input_channels - old_conv.in_channels\n            new_conv = nn.Conv2d(\n                in_channels=new_conv_channels, **old_conv_args)\n            model.backbone.conv1 = nn.Sequential(\n                # split input along channel dim\n                SplitTensor((old_conv.in_channels, new_conv_channels), dim=1),\n                # each split goes to its respective conv layer\n                Parallel(old_conv, new_conv),\n                # sum the parallel outputs\n                AddTensors())\n        elif input_channels < old_conv.in_channels:\n            model.backbone.conv1 = nn.Conv2d(\n                in_channels=input_channels, **old_conv_args)\n            model.backbone.conv1.weight.data[:, :input_channels] = \\\n                old_conv.weight.data[:, :input_channels]\n        else:\n            raise ConfigError(f'Something went wrong')\n\n        return model\n\n    def build_loss(self):\n        args = {}\n\n        loss_weights = self.cfg.solver.class_loss_weights\n        if loss_weights is not None:\n            loss_weights = torch.tensor(loss_weights, device=self.device)\n            args.update({'weight': loss_weights})\n\n        if self.cfg.solver.ignore_last_class:\n            num_classes = len(self.cfg.data.class_names)\n            args.update({'ignore_index': num_classes - 1})\n\n        loss = nn.CrossEntropyLoss(**args)\n\n        return loss\n\n    def train_step(self, batch, batch_ind):\n        x, y = batch\n        out = self.post_forward(self.model(x))\n        return {'train_loss': self.loss(out, y)}\n\n    def validate_step(self, batch, batch_ind):\n        x, y = batch\n        out = self.post_forward(self.model(x))\n        val_loss = self.loss(out, y)\n\n        num_labels = len(self.cfg.data.class_names)\n        y = y.view(-1)\n        out = self.prob_to_pred(out).view(-1)\n        conf_mat = compute_conf_mat(out, y, num_labels)\n\n        return {'val_loss': val_loss, 'conf_mat': conf_mat}\n\n    def validate_end(self, outputs, num_samples):\n        conf_mat = sum([o['conf_mat'] for o in outputs])\n        val_loss = torch.stack([o['val_loss']\n                                for o in outputs]).sum() / num_samples\n        conf_mat_metrics = compute_conf_mat_metrics(conf_mat,\n                                                    self.cfg.data.class_names)\n\n        metrics = {'val_loss': val_loss.item()}\n        metrics.update(conf_mat_metrics)\n\n        return metrics\n\n    def post_forward(self, x):\n        if isinstance(x, dict):\n            return x['out']\n        return x\n\n    def predict(self, x: torch.Tensor, raw_out: bool = False) -> torch.Tensor:\n        x = self.to_batch(x).float()\n        x = self.to_device(x, self.device)\n        with torch.no_grad():\n            out = self.model(x)\n            out = self.post_forward(out)\n            out = out.softmax(dim=1)\n            if not raw_out:\n                out = self.prob_to_pred(out)\n        out = self.to_device(out, 'cpu')\n        return out\n\n    def numpy_predict(self, x: np.ndarray,\n                      raw_out: bool = False) -> np.ndarray:\n        _, h, w, _ = x.shape\n        transform, _ = self.get_data_transforms()\n        x = self.normalize_input(x)\n        x = self.to_batch(x)\n        x = np.stack([transform(image=img)['image'] for img in x])\n        x = torch.from_numpy(x)\n        x = x.permute((0, 3, 1, 2))\n        out = self.predict(x, raw_out=True)\n        out = F.interpolate(\n            out, size=(h, w), mode='bilinear', align_corners=False)\n        out = self.prob_to_pred(out)\n        return self.output_to_numpy(out)\n\n    def prob_to_pred(self, x):\n        return x.argmax(1)\n\n    def plot_batch(self,\n                   x: torch.Tensor,\n                   y: Union[torch.Tensor, np.ndarray],\n                   output_path: str,\n                   z: Optional[torch.Tensor] = None,\n                   batch_limit: Optional[int] = None) -> None:\n        \"\"\"Plot a whole batch in a grid using plot_xyz.\n\n        Args:\n            x: batch of images\n            y: ground truth labels\n            output_path: local path where to save plot image\n            z: optional predicted labels\n            batch_limit: optional limit on (rendered) batch size\n        \"\"\"\n        batch_sz, c, h, w = x.shape\n        batch_sz = min(batch_sz,\n                       batch_limit) if batch_limit is not None else batch_sz\n        if batch_sz == 0:\n            return\n\n        channel_groups = self.cfg.data.channel_display_groups\n\n        nrows = batch_sz\n        # one col for each group + 1 for labels + 1 for predictions\n        ncols = len(channel_groups) + 1\n        if z is not None:\n            ncols += 1\n\n        fig, axes = plt.subplots(\n            nrows=nrows,\n            ncols=ncols,\n            squeeze=False,\n            constrained_layout=True,\n            figsize=(3 * ncols, 3 * nrows))\n\n        assert axes.shape == (nrows, ncols)\n\n        # (N, c, h, w) --> (N, h, w, c)\n        x = x.permute(0, 2, 3, 1)\n\n        # apply transform, if given\n        if self.cfg.data.plot_options.transform is not None:\n            tf = A.from_dict(self.cfg.data.plot_options.transform)\n            imgs = [tf(image=img)['image'] for img in x.numpy()]\n            x = torch.from_numpy(np.stack(imgs))\n\n        for i in range(batch_sz):\n            ax = (fig, axes[i])\n            if z is None:\n                self.plot_xyz(ax, x[i], y[i])\n            else:\n                self.plot_xyz(ax, x[i], y[i], z=z[i])\n\n        make_dir(output_path, use_dirname=True)\n        plt.savefig(output_path, bbox_inches='tight')\n        plt.close()\n\n    def plot_xyz(self,\n                 ax: Iterable,\n                 x: torch.Tensor,\n                 y: Union[torch.Tensor, np.ndarray],\n                 z: Optional[torch.Tensor] = None) -> None:\n\n        channel_groups = self.cfg.data.channel_display_groups\n\n        # make subplot titles\n        if not isinstance(channel_groups, dict):\n            channel_groups = {\n                f'Channels: {[*chs]}': chs\n                for chs in channel_groups\n            }\n\n        fig, ax = ax\n        img_axes = ax[:len(channel_groups)]\n        label_ax = ax[len(channel_groups)]\n\n        # plot input image(s)\n        for (title, chs), ch_ax in zip(channel_groups.items(), img_axes):\n            im = x[..., chs]\n            if len(chs) == 1:\n                # repeat single channel 3 times\n                im = im.expand(-1, -1, 3)\n            elif len(chs) == 2:\n                # add a 3rd channel with all pixels set to 0.5\n                h, w, _ = x.shape\n                third_channel = torch.full((h, w, 1), fill_value=.5)\n                im = torch.cat((im, third_channel), dim=-1)\n            elif len(chs) > 3:\n                # only use the first 3 channels\n                log.warn(f'Only plotting first 3 channels of channel-group '\n                         f'{title}: {chs}.')\n                im = x[..., chs[:3]]\n            ch_ax.imshow(im)\n            ch_ax.set_title(title)\n            ch_ax.set_xticks([])\n            ch_ax.set_yticks([])\n\n        class_colors = self.cfg.data.class_colors\n        colors = [color_to_triple(c) for c in class_colors]\n        colors = np.array(colors) / 255.\n        cmap = matplotlib.colors.ListedColormap(colors)\n\n        # plot labels\n        label_ax.imshow(\n            y, vmin=0, vmax=len(colors), cmap=cmap, interpolation='none')\n        label_ax.set_title(f'Ground truth labels')\n        label_ax.set_xticks([])\n        label_ax.set_yticks([])\n\n        # plot predictions\n        if z is not None:\n            pred_ax = ax[-1]\n            pred_ax.imshow(\n                z, vmin=0, vmax=len(colors), cmap=cmap, interpolation='none')\n            pred_ax.set_title(f'Predicted labels')\n            pred_ax.set_xticks([])\n            pred_ax.set_yticks([])\n\n        # add a legend to the rightmost subplot\n        class_names = self.cfg.data.class_names\n        legend_items = [\n            mpatches.Patch(facecolor=col, edgecolor='black', label=name)\n            for col, name in zip(colors, class_names)\n        ]\n        ax[-1].legend(\n            handles=legend_items,\n            loc='center right',\n            bbox_to_anchor=(1.8, 0.5))\n", "path": "rastervision_pytorch_learner/rastervision/pytorch_learner/semantic_segmentation_learner.py"}], "after_files": [{"content": "import warnings\nwarnings.filterwarnings('ignore')  # noqa\n\nfrom typing import Union, Iterable, Optional\n\nimport logging\n\nimport numpy as np\nimport matplotlib\nfrom matplotlib import pyplot as plt\nimport matplotlib.patches as mpatches\nmatplotlib.use('Agg')  # noqa\nimport albumentations as A\n\nimport torch\nfrom torch import nn\nfrom torch.nn import functional as F\nfrom torchvision import models\n\nfrom rastervision.pipeline.config import ConfigError\nfrom rastervision.pytorch_learner.learner import Learner\nfrom rastervision.pytorch_learner.utils import (\n    compute_conf_mat_metrics, compute_conf_mat, color_to_triple, SplitTensor,\n    Parallel, AddTensors)\nfrom rastervision.pipeline.file_system import make_dir\n\nlog = logging.getLogger(__name__)\n\n\nclass SemanticSegmentationLearner(Learner):\n    def build_model(self) -> nn.Module:\n        # TODO support FCN option\n        pretrained = self.cfg.model.pretrained\n        out_classes = len(self.cfg.data.class_names)\n        if self.cfg.solver.ignore_last_class:\n            out_classes -= 1\n        model = models.segmentation.segmentation._segm_model(\n            'deeplabv3',\n            self.cfg.model.get_backbone_str(),\n            out_classes,\n            False,\n            pretrained_backbone=pretrained)\n\n        input_channels = self.cfg.data.img_channels\n        old_conv = model.backbone.conv1\n\n        if input_channels == old_conv.in_channels:\n            return model\n\n        # these parameters will be the same for the new conv layer\n        old_conv_args = {\n            'out_channels': old_conv.out_channels,\n            'kernel_size': old_conv.kernel_size,\n            'stride': old_conv.stride,\n            'padding': old_conv.padding,\n            'dilation': old_conv.dilation,\n            'groups': old_conv.groups,\n            'bias': old_conv.bias\n        }\n\n        if not pretrained:\n            # simply replace the first conv layer with one with the\n            # correct number of input channels\n            new_conv = nn.Conv2d(in_channels=input_channels, **old_conv_args)\n            model.backbone.conv1 = new_conv\n            return model\n\n        if input_channels > old_conv.in_channels:\n            # insert a new conv layer parallel to the existing one\n            # and sum their outputs\n            new_conv_channels = input_channels - old_conv.in_channels\n            new_conv = nn.Conv2d(\n                in_channels=new_conv_channels, **old_conv_args)\n            model.backbone.conv1 = nn.Sequential(\n                # split input along channel dim\n                SplitTensor((old_conv.in_channels, new_conv_channels), dim=1),\n                # each split goes to its respective conv layer\n                Parallel(old_conv, new_conv),\n                # sum the parallel outputs\n                AddTensors())\n        elif input_channels < old_conv.in_channels:\n            model.backbone.conv1 = nn.Conv2d(\n                in_channels=input_channels, **old_conv_args)\n            model.backbone.conv1.weight.data[:, :input_channels] = \\\n                old_conv.weight.data[:, :input_channels]\n        else:\n            raise ConfigError(f'Something went wrong')\n\n        return model\n\n    def build_loss(self):\n        args = {}\n\n        loss_weights = self.cfg.solver.class_loss_weights\n        if loss_weights is not None:\n            loss_weights = torch.tensor(loss_weights, device=self.device)\n            args.update({'weight': loss_weights})\n\n        if self.cfg.solver.ignore_last_class:\n            num_classes = len(self.cfg.data.class_names)\n            args.update({'ignore_index': num_classes - 1})\n\n        loss = nn.CrossEntropyLoss(**args)\n\n        return loss\n\n    def train_step(self, batch, batch_ind):\n        x, y = batch\n        out = self.post_forward(self.model(x))\n        return {'train_loss': self.loss(out, y)}\n\n    def validate_step(self, batch, batch_ind):\n        x, y = batch\n        out = self.post_forward(self.model(x))\n        val_loss = self.loss(out, y)\n\n        num_labels = len(self.cfg.data.class_names)\n        y = y.view(-1)\n        out = self.prob_to_pred(out).view(-1)\n        conf_mat = compute_conf_mat(out, y, num_labels)\n\n        return {'val_loss': val_loss, 'conf_mat': conf_mat}\n\n    def validate_end(self, outputs, num_samples):\n        conf_mat = sum([o['conf_mat'] for o in outputs])\n        val_loss = torch.stack([o['val_loss']\n                                for o in outputs]).sum() / num_samples\n        conf_mat_metrics = compute_conf_mat_metrics(conf_mat,\n                                                    self.cfg.data.class_names)\n\n        metrics = {'val_loss': val_loss.item()}\n        metrics.update(conf_mat_metrics)\n\n        return metrics\n\n    def post_forward(self, x):\n        if isinstance(x, dict):\n            return x['out']\n        return x\n\n    def predict(self, x: torch.Tensor, raw_out: bool = False) -> torch.Tensor:\n        x = self.to_batch(x).float()\n        x = self.to_device(x, self.device)\n        with torch.no_grad():\n            out = self.model(x)\n            out = self.post_forward(out)\n            out = out.softmax(dim=1)\n            if not raw_out:\n                out = self.prob_to_pred(out)\n        out = self.to_device(out, 'cpu')\n        return out\n\n    def numpy_predict(self, x: np.ndarray,\n                      raw_out: bool = False) -> np.ndarray:\n        _, h, w, _ = x.shape\n        transform, _ = self.get_data_transforms()\n        x = self.normalize_input(x)\n        x = self.to_batch(x)\n        x = np.stack([transform(image=img)['image'] for img in x])\n        x = torch.from_numpy(x)\n        x = x.permute((0, 3, 1, 2))\n        out = self.predict(x, raw_out=True)\n        out = F.interpolate(\n            out, size=(h, w), mode='bilinear', align_corners=False)\n        out = self.prob_to_pred(out)\n        return self.output_to_numpy(out)\n\n    def prob_to_pred(self, x):\n        return x.argmax(1)\n\n    def plot_batch(self,\n                   x: torch.Tensor,\n                   y: Union[torch.Tensor, np.ndarray],\n                   output_path: str,\n                   z: Optional[torch.Tensor] = None,\n                   batch_limit: Optional[int] = None) -> None:\n        \"\"\"Plot a whole batch in a grid using plot_xyz.\n\n        Args:\n            x: batch of images\n            y: ground truth labels\n            output_path: local path where to save plot image\n            z: optional predicted labels\n            batch_limit: optional limit on (rendered) batch size\n        \"\"\"\n        batch_sz, c, h, w = x.shape\n        batch_sz = min(batch_sz,\n                       batch_limit) if batch_limit is not None else batch_sz\n        if batch_sz == 0:\n            return\n\n        channel_groups = self.cfg.data.channel_display_groups\n\n        nrows = batch_sz\n        # one col for each group + 1 for labels + 1 for predictions\n        ncols = len(channel_groups) + 1\n        if z is not None:\n            ncols += 1\n\n        fig, axes = plt.subplots(\n            nrows=nrows,\n            ncols=ncols,\n            squeeze=False,\n            constrained_layout=True,\n            figsize=(3 * ncols, 3 * nrows))\n\n        assert axes.shape == (nrows, ncols)\n\n        # (N, c, h, w) --> (N, h, w, c)\n        x = x.permute(0, 2, 3, 1)\n\n        # apply transform, if given\n        if self.cfg.data.plot_options.transform is not None:\n            tf = A.from_dict(self.cfg.data.plot_options.transform)\n            imgs = [tf(image=img)['image'] for img in x.numpy()]\n            x = torch.from_numpy(np.stack(imgs))\n\n        for i in range(batch_sz):\n            ax = (fig, axes[i])\n            if z is None:\n                self.plot_xyz(ax, x[i], y[i])\n            else:\n                self.plot_xyz(ax, x[i], y[i], z=z[i])\n\n        make_dir(output_path, use_dirname=True)\n        plt.savefig(output_path, bbox_inches='tight')\n        plt.close()\n\n    def plot_xyz(self,\n                 ax: Iterable,\n                 x: torch.Tensor,\n                 y: Union[torch.Tensor, np.ndarray],\n                 z: Optional[torch.Tensor] = None) -> None:\n\n        channel_groups = self.cfg.data.channel_display_groups\n\n        # make subplot titles\n        if not isinstance(channel_groups, dict):\n            channel_groups = {\n                f'Channels: {[*chs]}': chs\n                for chs in channel_groups\n            }\n\n        fig, ax = ax\n        img_axes = ax[:len(channel_groups)]\n        label_ax = ax[len(channel_groups)]\n\n        # plot input image(s)\n        for (title, chs), ch_ax in zip(channel_groups.items(), img_axes):\n            im = x[..., chs]\n            if len(chs) == 1:\n                # repeat single channel 3 times\n                im = im.expand(-1, -1, 3)\n            elif len(chs) == 2:\n                # add a 3rd channel with all pixels set to 0.5\n                h, w, _ = x.shape\n                third_channel = torch.full((h, w, 1), fill_value=.5)\n                im = torch.cat((im, third_channel), dim=-1)\n            elif len(chs) > 3:\n                # only use the first 3 channels\n                log.warn(f'Only plotting first 3 channels of channel-group '\n                         f'{title}: {chs}.')\n                im = x[..., chs[:3]]\n            ch_ax.imshow(im)\n            ch_ax.set_title(title)\n            ch_ax.set_xticks([])\n            ch_ax.set_yticks([])\n\n        class_colors = self.cfg.data.class_colors\n        colors = [color_to_triple(c) for c in class_colors]\n        colors = np.array(colors) / 255.\n        cmap = matplotlib.colors.ListedColormap(colors)\n\n        # plot labels\n        label_ax.imshow(\n            y, vmin=0, vmax=len(colors), cmap=cmap, interpolation='none')\n        label_ax.set_title(f'Ground truth labels')\n        label_ax.set_xticks([])\n        label_ax.set_yticks([])\n\n        # plot predictions\n        if z is not None:\n            pred_ax = ax[-1]\n            pred_ax.imshow(\n                z, vmin=0, vmax=len(colors), cmap=cmap, interpolation='none')\n            pred_ax.set_title(f'Predicted labels')\n            pred_ax.set_xticks([])\n            pred_ax.set_yticks([])\n\n        # add a legend to the rightmost subplot\n        class_names = self.cfg.data.class_names\n        legend_items = [\n            mpatches.Patch(facecolor=col, edgecolor='black', label=name)\n            for col, name in zip(colors, class_names)\n        ]\n        ax[-1].legend(\n            handles=legend_items,\n            loc='center right',\n            bbox_to_anchor=(1.8, 0.5))\n", "path": "rastervision_pytorch_learner/rastervision/pytorch_learner/semantic_segmentation_learner.py"}]}