From 4002c0d4ce4db3c9d8acf23138c3675b1c5c395f Mon Sep 17 00:00:00 2001 From: Ross Wightman Date: Sun, 1 Sep 2019 22:07:45 -0700 Subject: [PATCH] Fix AutoAugment abs translate calc --- timm/data/transforms.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/timm/data/transforms.py b/timm/data/transforms.py index d4f67bf9..33911638 100644 --- a/timm/data/transforms.py +++ b/timm/data/transforms.py @@ -179,8 +179,12 @@ def transforms_imagenet_train( transforms.RandomHorizontalFlip() ] if auto_augment: + if isinstance(img_size, tuple): + img_size_min = min(img_size) + else: + img_size_min = img_size aa_params = dict( - translate_const=img_size[-1] // 2 - 1, + translate_const=int(img_size_min * 0.45), img_mean=tuple([min(255, round(255 * x)) for x in mean]), ) if interpolation and interpolation != 'random':