rot_rand in SWIN transformer pre-training

[fengling0410]

Hi all, I'm confused about the rot_rand part of the Swin transformer pre-training when trying to reproduce the result. The rot_rand code is as the following:

def rot_rand(args, x_s):
    img_n = x_s.size()[0]
    x_aug = x_s.detach().clone()
    device = torch.device(f"cuda:{args.local_rank}")
    x_rot = torch.zeros(img_n).long().to(device) 
    for i in range(img_n):
        x = x_s[i]
        orientation = np.random.randint(0, 4)
        if orientation == 0:
            pass
        elif orientation == 1:
            x = x.rot90(1, (2, 3))
        elif orientation == 2:
            x = x.rot90(2, (2, 3)) 
        elif orientation == 3:
            x = x.rot90(3, (2, 3))
        x_aug[i] = x
        x_rot[i] = orientation
    return x_aug, x_rot

As far as I understand, x_s has the shape of (batch_size, no_channels, H, W, D), and in x = x_s[i] x has the shape of (no_channels, H, W, D). In order to rotate along the z-axis (same as rotate in x-y plane), the code should be x = x.rot90(1, (1, 2)) where 1 stands for H and 2 stands for W. I'm a little confused about why x = x.rot90(1, (2, 3)) is used here. May I ask for some advices regarding this? Thank you in advance :)

[tangy5]

Hi Fengling0410,
Thanks for your message. The axis tuple for rot90 is the rotation axis, and you are right, if your data is in [B, C, H, W, D], the rotation z axis will be set dimension [1,2] after extracting data point B. To our experience, the rotation prediction is an easier task as pre-texts. You can enable all 9 directions, it will be better.
Thanks.

[fengling0410]

Hi @tangy5, thank you for your reply. May I ask about the effect of rot_rand function in MONAI Swin transformer pre-training? Is it doing rotation along the z-axis or the x-axis? I think the Swin UNETR paper has made use of rotation along the z-axis. I'm going to pre-train the SWIN transformer based on the same dataset and then apply it to internal dataset, therefore I would like to dive deeper into the pre-training script. Thank you in advance :)

[tangy5]

We did it in z-axis. The implementation is bit modified after public release. If you are using the same data set, feel free to rotate along z-axis and [1,2] as rot90 rotation axis. Glad to see you will apply it to your dataset, and thanks again for your interests of the work. Wish you best of luck! Thank you!

[fengling0410]

Hi @tangy5, thank you again for your quick reply. Actually I've tried to use x = x.rot90(1, (1, 2)) to pre-train the model with other codes kept all the same. But the training plot and validation image on tensor board seems a little bit strange. I'm quite confused about this. Do you have any advices regarding this? Is it related to the modified implement you mentioned above? Here are the training plots I got:

The model is trained on single GPU with the command python main.py --use_checkpoint --batch_size=1 --num_steps=10000 --lrdecay --eval_num=100 --logdir='./runs/test' --lr=4e-4 --roi_x=96 --roi_y=96 --roi_z=96.