4. 常用环境变量

4.1. PYTORCH_DEFAULT_NCHW

对于四维Tensor，为获得更好的性能，mcPyTorch在其从CPU向device侧搬运时会更改layout为torch.channels_last (NHWC)。该环境变量用于恢复官方默认行为，即搬运时不做额外的layout改变逻辑。

备注

偶然情况下，上述默认更改layout的转换可能因为viewop的使用引起以下报错：

Runtime Error: view size is not compatible with input tensor's size and stride(...). Use .reshape(...) instead.

这时可以参考提示信息使用reshape替换viewop；也可设置 export PYTORCH_DEFAULT_NCHW=1 恢复官方默认行为，避免这种报错。

示例：

设置 unset PYTORCH_DEFAULT_NCHW，运行程序结果如下：

>>> import torch
>>> a = torch.rand(2,3,4,5).cuda()
>>> a.shape
torch.Size([2, 3, 4, 5])
>>> a.stride()
(60, 1, 15, 3)

设置 export PYTORCH_DEFAULT_NCHW=1，运行程序结果如下：

>>> import torch
>>> a = torch.rand(2,3,4,5).cuda()
>>> a.shape
torch.Size([2, 3, 4, 5])
>>> a.stride()
(60, 20, 5, 1)

torch.backends.cuda.matmul.allow_tf32 默认为 False。设置该环境变量时控制 torch.backends.cuda.matmul.allow_tf32 默认为 True。