This is a custom version of DataParallel that works better with our training data. It should also be faster than the general case.
| 151 | |
| 152 | |
| 153 | class CustomDataParallel(nn.DataParallel): |
| 154 | """ |
| 155 | This is a custom version of DataParallel that works better with our training data. |
| 156 | It should also be faster than the general case. |
| 157 | |
| 158 | """ |
| 159 | def scatter(self, inputs, kwargs, device_ids): |
| 160 | # More like scatter and data prep at the same time. The point is we prep the data in such a way |
| 161 | # that no scatter is necessary, and there's no need to shuffle stuff around different GPUs. |
| 162 | devices = ['cuda:' + str(x) for x in device_ids] |
| 163 | splits = self.prepare_data(inputs[0], devices, allocation=args.batch_alloc) |
| 164 | |
| 165 | return [[split[device_idx] for split in splits] for device_idx in range(len(devices))], \ |
| 166 | [kwargs] * len(devices) |
| 167 | |
| 168 | def gather(self, outputs, output_device): |
| 169 | out = {} |
| 170 | |
| 171 | for k in outputs[0]: |
| 172 | out[k] = torch.stack([output[k].to(output_device) for output in outputs]) |
| 173 | |
| 174 | return out |
| 175 | |
| 176 | @torch.no_grad() |
| 177 | def prepare_data(self, datum, devices:list=None, allocation:list=None): |
| 178 | |
| 179 | def gradinator(x): |
| 180 | x.requires_grad = False |
| 181 | return x |
| 182 | if devices is None: |
| 183 | devices = ['cuda:0'] |
| 184 | if allocation is None: |
| 185 | allocation = [args.batch_size // len(devices)] * (len(devices) - 1) |
| 186 | allocation.append(args.batch_size - sum(allocation)) # The rest might need more/less |
| 187 | |
| 188 | batched_images, batched_gt_instances, batched_gt_depths = datum |
| 189 | |
| 190 | cur_idx = 0 |
| 191 | for device, alloc in zip(devices, allocation): |
| 192 | for _ in range(alloc): |
| 193 | batched_images[cur_idx] = gradinator(batched_images[cur_idx].to(device)) |
| 194 | batched_gt_depths[cur_idx] = gradinator(batched_gt_depths[cur_idx].to(device)) |
| 195 | for key in batched_gt_instances[cur_idx]: |
| 196 | batched_gt_instances[cur_idx][key] = gradinator(batched_gt_instances[cur_idx][key].to(device)) |
| 197 | cur_idx += 1 |
| 198 | |
| 199 | if cfg.preserve_aspect_ratio: |
| 200 | # Choose a random size from the batch |
| 201 | _, h, w = batched_images[random.randint(0, len(batched_images)-1)].size() |
| 202 | for idx, (image, gt_depth, gt_instances) in enumerate(zip(batched_images, batched_gt_depths, batched_gt_instances)): |
| 203 | batched_images[idx], batched_gt_depths[idx], batched_gt_instances[idx] \ |
| 204 | = enforce_size(image, gt_depth, gt_instances, w, h) |
| 205 | cur_idx = 0 |
| 206 | split_images, split_depths, split_instances = [[None for alloc in allocation] for _ in range(3)] |
| 207 | |
| 208 | for device_idx, alloc in enumerate(allocation): |
| 209 | split_images[device_idx] = torch.stack(batched_images[cur_idx:cur_idx+alloc], dim=0) |
| 210 | split_depths[device_idx] = torch.stack(batched_gt_depths[cur_idx:cur_idx+alloc], dim=0) |