Pytorch

加载数据类

获取数据和其对应的label值、总共有多少数据

from torch.utils.data import Dataset

GPU训练

将输入、输出、模型、损失函数都移到cuda上

设置环境变量

必须在import torch之前键入以下代码

import os
os.environ['CUDA_VISIBLE_DEVICES'] = '3'

然后都移动到cuda上

x = x.cuda()
y = y.cuda()
model = model.cuda()
loss_function = loss_function.cuda()

使用todevice

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

model.to(device)

Tensorboard

from torch.utils.tensorboard import SummaryWriter


writer = SummaryWriter("logs")
# y = x+1
for i in range(100):
    writer.add_scalar("y=x",i+1,i)
writer.close()

生成一个logs文件夹，终端进入logs文件夹的同级目录，输入

tensroboard --logdir=logs --port 6007

计算模型体量

进度条

from tqdm import tqdm  # 直接import tqdm会报错

loop = tqdm(enumerate(train_dataloader), total=train_data_size)
for step, data in loop:
  pass

保存模型

包含GPU tensors的模型，自动load到GPU上

保存整个模型

## 保存模型
torch.save(model, MODEL_PATH)

## 读取模型
model = torch.load(MODEL_PATH)

只保存参数

## 保存模型
torch.save(model.state_dict(), MODEL_PATH)
torch.save(optimizer.state_dict(), OPTIM_PATH)

## 读取模型
if os.path.exists(MODEL_PATH)
    model.load_state_dict(torch.load(MODEL_PATH))
if os.path.exists(OPTIM_PATH)
  optimizer.load_state_dict(optimizer.state_dict(), OPTIM_PATH)