| 568 | return self.beta_func(self.beta, self.beta, size=(batch,)) |
| 569 | |
| 570 | def forward(self, data1, label1, data2, label2): |
| 571 | assert all( |
| 572 | isinstance(inp, Tensor) for inp in [data1, label1, data2, label2] |
| 573 | ), "expected input is megengine.Tensor" |
| 574 | |
| 575 | batch, C, H, W = data1.shape |
| 576 | self.lamb = self.sample(batch) |
| 577 | |
| 578 | label = self.lamb * label1 + (1.0 - self.lamb) * label2 |
| 579 | |
| 580 | data = ( |
| 581 | self.lamb.reshape(batch, 1, 1, 1) * data1 |
| 582 | + (1 - self.lamb).reshape(batch, 1, 1, 1) * data2 |
| 583 | ) |
| 584 | |
| 585 | return data, label |
| 586 | |
| 587 | |
| 588 | class Cutmix(Module): |