1-torch常用操作和函数

1. 基础操作

1.1 导入包和版本查询

import torch
import torch.nn as nn
import torchvision
print(torch.__version__) # 查看toch版本
print(torch.version.cuda) # 查看cuda版本
print(torch.backends.cudnn.version()) # 查看cudnn版本
print(torch.cuda.get_device_name(0)) #查看显卡名字

1.2 固定随机种子

def seed_torch(seed=42):
    random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.backends.cudnn.deterministic = True

seed_torch(seed=CFG.seed)

1.3 显卡设置

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') #指定是cuda:2

# 多张显卡
import os
os.environ['CUDA_VISIBLE_DEVICES'] = '1,3'
# 另一种方式
CUDA_VISIBLE_DEVICES=0,1 python train.py

1.4 显示显存占用-清除缓存

mem = torch.cuda.memory_reserved() / 1E9 if torch.cuda.is_available() else 0
torch.cuda.empty_cache()

1.5 查询tensor基本信息

tensor = torch.randn(3,4,5)
print(tensor.type())  # 数据类型 输出：torch.FloatTensor
print(tensor.size())  # 张量的shape，是个元组输出：torch.Size([3, 4, 5])
print(tensor.dim())   # 维度的数量，输出：3

1.6 tensor数据类型转换

# 设置默认类型，pytorch中的FloatTensor远远快于DoubleTensor
torch.set_default_tensor_type(torch.FloatTensor)
# 基础类型转换
tensor = tensor.cuda()
tensor = tensor.cpu()
tensor = tensor.float()
tensor = tensor.long()

和numpy的相互转换

ndarray = tensor.cpu().numpy() #tensor代表具体得变量，ndarray具体变量
tensor = torch.from_numpy(ndarray).float()
tensor = torch.from_numpy(ndarray.copy()).float() # If ndarray has negative stride.

sample = cv2.imread('R.jpg')[:,:,::-1]
torch.from_numpy(sample.copy()).unsqueeze(0).permute(0, 3, 1, 2)

和tensor和PIL转换

image = PIL.Image.fromarray(torch.clamp(tensor*255, min=0, max=255).byte().permute(1,2,0).cpu().numpy())

path = r'./test.jpg'
tensor = torch.from_numpy(np.asarray(PIL.Image.open(path))).permute(2,0,1).float() / 255

numpy和pil的相互转换

image = PIL.Image.fromarray(ndarray.astype(np.uint8))
ndarray = np.asarray(PIL.Image.open(path))

1.7 拼接张量

'''
注意torch.cat和torch.stack的区别在于torch.cat沿着给定的维度拼接，
而torch.stack会新增一维。例如当参数是3个10x5的张量，torch.cat的结果是30x5的张量，
而torch.stack的结果是3x10x5的张量。
'''
tensor = torch.cat(list_of_tensors, dim=0)
tensor = torch.stack(list_of_tensors, dim=0)

1.8 标签转one_hot编码

# pytorch的标记默认从0开始
tensor = torch.tensor([0, 2, 1, 3])
N = tensor.size(0)
num_classes = 4
one_hot = torch.zeros(N, num_classes).long()
one_hot.scatter_(dim=1, index=torch.unsqueeze(tensor, dim=1), src=torch.ones(N, num_classes).long())

1.9张量扩展

# Expand tensor of shape 64*512 to shape 64*512*7*7.
tensor = torch.rand(64,512)
torch.reshape(tensor, (64, 512, 1, 1)).expand(64, 512, 7, 7)

2. 简单网络定义

# convolutional neural network (2 convolutional layers)
class ConvNet(nn.Module):
    def __init__(self, num_classes=10):
        super(ConvNet, self).__init__()
        self.layer1 = nn.Sequential(
            nn.Conv2d(1, 16, kernel_size=5, stride=1, padding=2),
            nn.BatchNorm2d(16),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2))
        self.layer2 = nn.Sequential(
            nn.Conv2d(16, 32, kernel_size=5, stride=1, padding=2),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2))
        self.fc = nn.Linear(7*7*32, num_classes)
    def forward(self, x):
        out = self.layer1(x)
        out = self.layer2(out)
        out = out.reshape(out.size(0), -1)
        out = self.fc(out)
        return out
model = ConvNet(num_classes).to(device)

3.显示一个batch图片

def show_img(img, mask=None):
    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
#     img = clahe.apply(img)
#     plt.figure(figsize=(10,10))
    plt.imshow(img, cmap='bone')
    
    if mask is not None:
        # plt.imshow(np.ma.masked_where(mask!=1, mask), alpha=0.5, cmap='autumn')
        plt.imshow(mask, alpha=0.5)
        handles = [Rectangle((0,0),1,1, color=_c) for _c in [(0.667,0.0,0.0), (0.0,0.667,0.0), (0.0,0.0,0.667)]]
        labels = ["Large Bowel", "Small Bowel", "Stomach"]
        plt.legend(handles,labels)
    plt.axis('off')
    
def plot_batch(imgs, msks, size=3):
    plt.figure(figsize=(5*5, 5))
    for idx in range(size):
        plt.subplot(1, 5, idx+1)
        img = imgs[idx,].permute((1, 2, 0)).numpy()*255.0
        img = img.astype('uint8')
        msk = msks[idx,].permute((1, 2, 0)).numpy()*255.0
        show_img(img, msk)
    plt.tight_layout()
    plt.show()

imgs, msks = next(iter(train_loader))
plot_batch(imgs, msks, size=5)