(self, batch_size, is_distributed, no_aug=False)
| 79 | return self.model |
| 80 | |
| 81 | def get_data_loader(self, batch_size, is_distributed, no_aug=False): |
| 82 | from yolox.data import ( |
| 83 | COCODataset, |
| 84 | DataLoader, |
| 85 | InfiniteSampler, |
| 86 | MosaicDetection, |
| 87 | TrainTransform, |
| 88 | YoloBatchSampler |
| 89 | ) |
| 90 | |
| 91 | dataset = COCODataset( |
| 92 | data_dir=None, |
| 93 | json_file=self.train_ann, |
| 94 | img_size=self.input_size, |
| 95 | preproc=TrainTransform( |
| 96 | rgb_means=(0.485, 0.456, 0.406), |
| 97 | std=(0.229, 0.224, 0.225), |
| 98 | max_labels=50, |
| 99 | ), |
| 100 | ) |
| 101 | |
| 102 | dataset = MosaicDetection( |
| 103 | dataset, |
| 104 | mosaic=not no_aug, |
| 105 | img_size=self.input_size, |
| 106 | preproc=TrainTransform( |
| 107 | rgb_means=(0.485, 0.456, 0.406), |
| 108 | std=(0.229, 0.224, 0.225), |
| 109 | max_labels=120, |
| 110 | ), |
| 111 | degrees=self.degrees, |
| 112 | translate=self.translate, |
| 113 | scale=self.scale, |
| 114 | shear=self.shear, |
| 115 | perspective=self.perspective, |
| 116 | enable_mixup=self.enable_mixup, |
| 117 | ) |
| 118 | |
| 119 | self.dataset = dataset |
| 120 | |
| 121 | if is_distributed: |
| 122 | batch_size = batch_size // dist.get_world_size() |
| 123 | |
| 124 | sampler = InfiniteSampler(len(self.dataset), seed=self.seed if self.seed else 0) |
| 125 | |
| 126 | batch_sampler = YoloBatchSampler( |
| 127 | sampler=sampler, |
| 128 | batch_size=batch_size, |
| 129 | drop_last=False, |
| 130 | input_dimension=self.input_size, |
| 131 | mosaic=not no_aug, |
| 132 | ) |
| 133 | |
| 134 | dataloader_kwargs = {"num_workers": self.data_num_workers, "pin_memory": True} |
| 135 | dataloader_kwargs["batch_sampler"] = batch_sampler |
| 136 | train_loader = DataLoader(self.dataset, **dataloader_kwargs) |
| 137 | |
| 138 | return train_loader |
no test coverage detected