【PyTorch总结】NPU与GPU上数据类型转换

NPU上

a = torch.Tensor([1.]).npu() # tensor([1.], device='npu:0')
b = torch.Tensor([1.]).npu() # tensor([1.], device='npu:0')
c = a < b # tensor([False], device='npu:0')

### 下列操作支持
c = torch.tensor(c, dtype=a.dtype) # tensor([0.], device='npu:0')
c = c.to(a.dtype) # tensor([0.], device='npu:0')
c = c.type(a.dtype) # tensor([0.], device='npu:0')

### 下列操作不支持
c.astype(a.type()) # AttributeError: 'Tensor' object has no attribute 'astype'
c.type(a.type()) # ValueError: invalid type: 'torch.npu.FloatTensor'

GPU上

import torch
a = torch.Tensor([1.]).cuda()
b = torch.Tensor([1.]).cuda()

c = a < b
c = torch.tensor(c, dtype=a.dtype)
print(c) # tensor([0.], device='cuda:0')

c = a < b
c = c.to(a.dtype)
print(c) # tensor([0.], device='cuda:0')

c = a < b
c = c.type(a.dtype)
print(c) # tensor([0.], device='cuda:0')


c.type(a.type())
print(c) # tensor([0.], device='cuda:0')

# 注意pytorch中没有astype函数，下面这种做法是错误的。
c.astype(a.type())