Method prepare_data

train.py:177–213

(self, datum, devices:list=None, allocation:list=None)

    
    @torch.no_grad()
    def prepare_data(self, datum, devices:list=None, allocation:list=None):

        def gradinator(x):
            x.requires_grad = False
            return x
        if devices is None:
            devices = ['cuda:0']
        if allocation is None:
            allocation = [args.batch_size // len(devices)] * (len(devices) - 1)
            allocation.append(args.batch_size - sum(allocation)) # The rest might need more/less
        
        batched_images, batched_gt_instances, batched_gt_depths = datum

        cur_idx = 0
        for device, alloc in zip(devices, allocation):
            for _ in range(alloc):
                batched_images[cur_idx]  = gradinator(batched_images[cur_idx].to(device))
                batched_gt_depths[cur_idx]  = gradinator(batched_gt_depths[cur_idx].to(device))
                for key in batched_gt_instances[cur_idx]:
                    batched_gt_instances[cur_idx][key] = gradinator(batched_gt_instances[cur_idx][key].to(device))
                cur_idx += 1

        if cfg.preserve_aspect_ratio:
            # Choose a random size from the batch
            _, h, w = batched_images[random.randint(0, len(batched_images)-1)].size()
            for idx, (image, gt_depth, gt_instances) in enumerate(zip(batched_images, batched_gt_depths, batched_gt_instances)):
                batched_images[idx], batched_gt_depths[idx], batched_gt_instances[idx] \
                    = enforce_size(image, gt_depth, gt_instances, w, h)
        cur_idx = 0
        split_images, split_depths, split_instances = [[None for alloc in allocation] for _ in range(3)]

        for device_idx, alloc in enumerate(allocation):
            split_images[device_idx]    = torch.stack(batched_images[cur_idx:cur_idx+alloc], dim=0)
            split_depths[device_idx]    = torch.stack(batched_gt_depths[cur_idx:cur_idx+alloc], dim=0)
            split_instances[device_idx]   = batched_gt_instances[cur_idx:cur_idx+alloc]
            cur_idx += alloc
        return split_images, split_instances, split_depths


def train():