mirrors/pytorch-image-models
1
0
Fork 0
mirror of https://github.com/huggingface/pytorch-image-models.git synced 2025-06-03 15:01:08 +08:00
Code Issues Projects Releases Wiki Activity

Merge pull request #1222 from Leoooo333/master

Browse Source 
Fix mixup/one_hot device problem
This commit is contained in:
Ross Wightman 2023-05-10 08:59:15 -07:00 committed by GitHub
commit 8ce9a2c00a
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 7 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
timm/data/mixup.py
View File

@ -14,16 +14,16 @@ import numpy as np
import torch
def one_hot(x, num_classes, on_value=1., off_value=0., device='cuda'):
def one_hot(x, num_classes, on_value=1., off_value=0.):
x = x.long().view(-1, 1)
return torch.full((x.size()[0], num_classes), off_value, device=device).scatter_(1, x, on_value)
return torch.full((x.size()[0], num_classes), off_value, device=x.device).scatter_(1, x, on_value)
def mixup_target(target, num_classes, lam=1., smoothing=0.0, device='cuda'):
def mixup_target(target, num_classes, lam=1., smoothing=0.0):
off_value = smoothing / num_classes
on_value = 1. - smoothing + off_value
y1 = one_hot(target, num_classes, on_value=on_value, off_value=off_value, device=device)
y2 = one_hot(target.flip(0), num_classes, on_value=on_value, off_value=off_value, device=device)
y1 = one_hot(target, num_classes, on_value=on_value, off_value=off_value)
y2 = one_hot(target.flip(0), num_classes, on_value=on_value, off_value=off_value)
return y1 * lam + y2 * (1. - lam)
@ -214,7 +214,7 @@ class Mixup:
lam = self._mix_pair(x)
else:
lam = self._mix_batch(x)
target = mixup_target(target, self.num_classes, lam, self.label_smoothing, x.device)
target = mixup_target(target, self.num_classes, lam, self.label_smoothing)
return x, target
@ -310,7 +310,7 @@ class FastCollateMixup(Mixup):
else:
lam = self._mix_batch_collate(output, batch)
target = torch.tensor([b[1] for b in batch], dtype=torch.int64)
target = mixup_target(target, self.num_classes, lam, self.label_smoothing, device='cpu')
target = mixup_target(target, self.num_classes, lam, self.label_smoothing)
target = target[:batch_size]
return output, target