如何用PyTorch Torchvision实现仅45°/-45°的随机图像旋转

要实现仅以50%概率选择45°或-45°旋转图像，你可以用以下几种简洁方法替代原有的RandomRotation：

方法1：使用RandomChoice（最简便） #

torchvision.transforms.RandomChoice会从给定的变换列表中等概率随机选择一个执行，刚好匹配你50%概率选两种角度的需求：

import torchvision.transforms as T

transforms = T.Compose([
    T.RandomChoice([
        T.RandomRotation([45, 45]),  # 固定旋转45°
        T.RandomRotation([-45, -45]) # 固定旋转-45°
    ])
])

给RandomRotation传入长度为1的区间，就能确保每次只应用指定角度的旋转。

方法2：自定义变换类 #

如果需要后续调整概率等更灵活的逻辑，可以自定义一个继承自T.Transform的类：

import torchvision.transforms as T
import random

class FixedAngleRandomRotation(T.Transform):
    def __init__(self, angles):
        self.angles = angles
    
    def __call__(self, img):
        angle = random.choice(self.angles)
        return T.functional.rotate(img, angle)

# 使用时传入目标角度列表，默认等概率选择
transforms = T.Compose([
    FixedAngleRandomRotation([45, -45])
])

后续如果要调整两种角度的概率，只需把random.choice换成random.choices并指定权重即可。

方法3：使用Lambda变换 #

用T.Lambda快速封装简单的随机旋转逻辑，代码最简洁：

import torchvision.transforms as T
import random

transforms = T.Compose([
    T.Lambda(lambda img: T.functional.rotate(img, random.choice([45, -45])))
])

内容的提问来源于stack exchange，提问作者Lézard