This is a custom version of DataParallel that works better with our training data. It should also be faster than the general case.
| 147 | return losses |
| 148 | |
| 149 | class CustomDataParallel(nn.DataParallel): |
| 150 | """ |
| 151 | This is a custom version of DataParallel that works better with our training data. |
| 152 | It should also be faster than the general case. |
| 153 | """ |
| 154 | |
| 155 | def scatter(self, inputs, kwargs, device_ids): |
| 156 | # More like scatter and data prep at the same time. The point is we prep the data in such a way |
| 157 | # that no scatter is necessary, and there's no need to shuffle stuff around different GPUs. |
| 158 | devices = ['cuda:' + str(x) for x in device_ids] |
| 159 | splits = prepare_data(inputs[0], devices, allocation=args.batch_alloc) |
| 160 | |
| 161 | return [[split[device_idx] for split in splits] for device_idx in range(len(devices))], \ |
| 162 | [kwargs] * len(devices) |
| 163 | |
| 164 | def gather(self, outputs, output_device): |
| 165 | out = {} |
| 166 | |
| 167 | for k in outputs[0]: |
| 168 | out[k] = torch.stack([output[k].to(output_device) for output in outputs]) |
| 169 | |
| 170 | return out |
| 171 | |
| 172 | def train(): |
| 173 | if not os.path.exists(args.save_folder): |