pytorch_classification/deep_thoughts -bozhu/pytoch32 GPU/作业/gcn.py
单击单卡：
单机多卡：DataParallel
单机多卡DDP：
示例1：G:\deep-learn\deep-learning\pytorch_classification\train_multi_GPU\train_single_gpu.py
train_single_gpu.py：
G:\deep-learn\deep-learning\pytorch_classification\train_multi_GPU\train_multi_gpu_using_launch.py
示例1：PyTorch单卡/多卡下模型保存/加载
由于训练和测试所使用的硬件条件不同，在模型的保存和加载过程中可能因为单GPU和多GPU环境的不同带来模型不匹配等问题。这里对PyTorch框架下单卡/多卡下模型的保存和加载问题进行排列组合（=4），样例模型是torchvision中预训练模型resnet152，不尽之处欢迎大家补充。
1 数据格式与存储内容
1.1 模型存储格式
PyTorch存储模型主要采用pkl，pt，pth三种格式。就使用层面来说没有区别，这里不做具体的讨论。第5部分中列出了笔者查阅到的一些资料，感兴趣的读者可以进一步研究，欢迎留言讨论。
1.2 模型存储内容
一个PyTorch模型主要包含两个部分：模型结构和权重。其中模型是继承nn.Module的类，权重的数据结构是一个字典（key是层名，value是权重向量）。存储也由此分为两种形式：存储整个模型（包括结构和权重），和只存储模型权重。
from torchvision import models
model = models.resnet152(pretrained=True)
# 保存整个模型
torch.save(model, save_dir)
# 保存模型权重
torch.save(model.state_dict, save_dir)对于PyTorch而言，pt, pth和pkl三种数据格式均支持模型权重和整个模型的存储，因此使用上没有差别。
2 单卡和多卡模型存储的区别
PyTorch中将模型和数据放到GPU上有两种方式——.cuda和.to(device)，笔者一般使用前者。之后如果要使用多卡训练的话，需要对模型使用torch.nn.DataParallel。示例如下：
os.environ['CUDA_VISIBLE_DEVICES'] = '0' # 如果是多卡改成类似0,1,2
model = model.cuda() # 单卡
model = torch.nn.DataParallel(model).cuda() # 多卡之后我们把model对应的layer名称打印出来看一下，可以观察到差别在于多卡并行的模型每层的名称前多了一个“module”：
单卡层名：
多卡层名
这种模型表示的不同可能会导致模型保存和加载过程中需要处理一些conflicts，下面对各种 可能的情况做分类讨论。
3 情况分类讨论
3.1 单卡保存+单卡加载
在使用os.envision命令指定使用的GPU后，即可进行模型保存读取操作。注意这里即便保存和读取时使用的GPU不同也无妨。
import os
import torch
from torchvision import models
os.environ['CUDA_VISIBLE_DEVICES'] = '0' #这里替换成希望使用的GPU编号
model = models.resnet152(pretrained=True)
model.cuda()
# 保存+读取整个模型
torch.save(model, save_dir)
loaded_model = torch.load(save_dir)
loaded_model.cuda()
# 保存+读取模型权重
torch.save(model.state_dict(), save_dir)
loaded_dict = torch.load(save_dir)
loaded_model = models.resnet152() #注意这里需要对模型结构有定义
loaded_model.state_dict = loaded_dict
loaded_model.cuda()3.2 单卡保存+多卡加载
这种情况的处理比较简单，读取单卡保存的模型后，使用nn.DataParallel函数进行分布式训练设置即可（相当于3.1代码中.cuda()替换一下）：
import os
import torch
from torchvision import models
os.environ['CUDA_VISIBLE_DEVICES'] = '0' #这里替换成希望使用的GPU编号
model = models.resnet152(pretrained=True)
model.cuda()
# 保存+读取整个模型
torch.save(model, save_dir)
os.environ['CUDA_VISIBLE_DEVICES'] = '1,2' #这里替换成希望使用的GPU编号
loaded_model = torch.load(save_dir)
loaded_model = nn.DataParallel(loaded_model).cuda()
# 保存+读取模型权重
torch.save(model.state_dict(), save_dir)
os.environ['CUDA_VISIBLE_DEVICES'] = '1,2' #这里替换成希望使用的GPU编号
loaded_dict = torch.load(save_dir)
loaded_model = models.resnet152() #注意这里需要对模型结构有定义
loaded_model.state_dict = loaded_dict
loaded_model = nn.DataParallel(loaded_model).cuda()3.3 多卡保存+单卡加载
这种情况下的核心问题是：如何去掉权重字典键名中的"module"，以保证模型的统一性。
对于加载整个模型，直接提取模型的module属性即可：
import os
import torch
from torchvision import models
os.environ['CUDA_VISIBLE_DEVICES'] = '1,2' #这里替换成希望使用的GPU编号
model = models.resnet152(pretrained=True)
model = nn.DataParallel(model).cuda()
# 保存+读取整个模型
torch.save(model, save_dir)
os.environ['CUDA_VISIBLE_DEVICES'] = '0' #这里替换成希望使用的GPU编号
loaded_model = torch.load(save_dir)
loaded_model = loaded_model.module对于加载模型权重，有以下几种思路：
- （推荐）去除字典里的module麻烦，往model里添加module简单：
import os import torch from torchvision import models os.environ['CUDA_VISIBLE_DEVICES'] = '0,1,2' #这里替换成希望使用的GPU编号 model = models.resnet152(pretrained=True) model = nn.DataParallel(model).cuda() # 保存+读取模型权重 torch.save(model.state_dict(), save_dir) os.environ['CUDA_VISIBLE_DEVICES'] = '0' #这里替换成希望使用的GPU编号 loaded_dict = torch.load(save_dir) loaded_model = models.resnet152() #注意这里需要对模型结构有定义 loaded_model = nn.DataParallel(loaded_model).cuda() loaded_model.state_dict = loaded_dict这样即便是单卡，也可以开始训练了（相当于分布到单卡上）
- 遍历字典去除module
from collections import OrderedDict os.environ['CUDA_VISIBLE_DEVICES'] = '0' #这里替换成希望使用的GPU编号 loaded_dict = torch.load(save_dir) new_state_dict = OrderedDict() for k, v in loaded_dict.items(): name = k[7:] # module字段在最前面，从第7个字符开始就可以去掉module new_state_dict[name] = v #新字典的key值对应的value一一对应 loaded_model = models.resnet152() #注意这里需要对模型结构有定义 loaded_model.state_dict = new_state_dict loaded_model = loaded_model.cuda()3.4 多卡保存+多卡加载
由于是模型保存和加载都使用的是多卡，因此不存在模型层名前缀不同的问题。但多卡状态下存在一个device（使用的GPU）匹配的问题，即保存整个模型时会同时保存所使用的GPU id等信息，读取时若这些信息和当前使用的GPU信息不符则可能会报错或者程序不按预定状态运行。具体表现为以下两点：
- 读取整个模型再使用nn.DataParallel进行分布式训练设置
-
这种情况很可能会造成保存的整个模型中GPU id和读取环境下设置的GPU id不符，训练时数据所在device和模型所在device不一致而报错。
- 读取整个模型而不使用nn.DataParallel进行分布式训练设置
这种情况可能不会报错，测试中发现程序会自动使用设备的前n个GPU进行训练（n是保存的模型使用的GPU个数）。此时如果指定的GPU个数少于n，则会报错。在这种情况下，只有保存模型时环境的device id和读取模型时环境的device id一致，程序才会按照预期在指定的GPU上进行分布式训练。
相比之下，读取模型权重，之后再使用nn.DataParallel进行分布式训练设置则没有问题。因此多卡模式下建议使用权重的方式存储和读取模型：
import os
import torch
from torchvision import models
os.environ['CUDA_VISIBLE_DEVICES'] = '0,1,2' #这里替换成希望使用的GPU编号
model = models.resnet152(pretrained=True)
model = nn.DataParallel(model).cuda()
# 保存+读取模型权重，强烈建议！！
torch.save(model.state_dict(), save_dir)
loaded_dict = torch.load(save_dir)
loaded_model = models.resnet152() #注意这里需要对模型结构有定义
loaded_model = nn.DataParallel(loaded_model).cuda()
loaded_model.state_dict = loaded_dict如果只有保存的整个模型，也可以采用提取权重的方式构建新的模型：
# 读取整个模型
loaded_whole_model = torch.load(save_dir)
loaded_model = models.resnet152() #注意这里需要对模型结构有定义
loaded_model.state_dict = loaded_whole_model.state_dict
loaded_model = nn.DataParallel(loaded_model).cuda()------
另外，上面所有对于loaded_model修改权重字典的形式都是通过赋值来实现的，在PyTorch中还可以通过"load_state_dict"函数来实现：
loaded_model.load_state_dict(loaded_dict)示例2：跨设备保存加载模型
1、在 CPU 上加载在 GPU 上训练并保存的模型（Save on GPU, Load on CPU）：
device = torch.device('cpu')
model = TheModelClass()
# Load all tensors onto the CPU device
model.load_state_dict(torch.load('net_params.pkl', map_location=device))map_location：a function, torch.device, string or a dict specifying how to remap storage locations
令 torch.load() 函数的 map_location 参数等于 torch.device('cpu') 即可。 这里令 map_location 参数等于 'cpu' 也同样可以。
2、在 GPU 上加载在 GPU 上训练并保存的模型（Save on GPU, Load on GPU）：
device = torch.device("cuda")
model = TheModelClass()
model.load_state_dict(torch.load('net_params.pkl'))
model.to(device)在这里使用 map_location 参数不起作用，要使用 model.to(torch.device("cuda")) 将模型转换为CUDA优化的模型。
还需要对将要输入模型的数据调用 data = data.to(device)，即将数据从CPU转移到GPU。请注意，调用 my_tensor.to(device) 会返回一个 my_tensor 在 GPU 上的副本，它不会覆盖 my_tensor。因此需要手动覆盖张量：my_tensor = my_tensor.to(device)。
3、在 GPU 上加载在 GPU 上训练并保存的模型（Save on CPU, Load on GPU）
device = torch.device("cuda")
model = TheModelClass()
model.load_state_dict(torch.load('net_params.pkl', map_location="cuda:0"))
model.to(device)当加载包含GPU tensors的模型时，这些tensors 会被默认加载到GPU上，不过是同一个GPU设备。
当有多个GPU设备时，可以通过将 map_location 设定为 cuda:device_id 来指定使用哪一个GPU设备，上面例子是指定编号为0的GPU设备。
其实也可以将 torch.device("cuda") 改为 torch.device("cuda:0") 来指定编号为0的GPU设备。
最后调用 model.to(torch.device('cuda')) 来将模型的tensors转换为 CUDA tensors。
下面是PyTorch官方文档上的用法，可以进行参考：
>>> torch.load('tensors.pt')
# Load all tensors onto the CPU
>>> torch.load('tensors.pt', map_location=torch.device('cpu'))
# Load all tensors onto the CPU, using a function
>>> torch.load('tensors.pt', map_location=lambda storage, loc: storage)
# Load all tensors onto GPU 1
>>> torch.load('tensors.pt', map_location=lambda storage, loc: storage.cuda(1))
# Map tensors from GPU 1 to GPU 0
>>> torch.load('tensors.pt', map_location={'cuda:1':'cuda:0'})有时我们需要把数据和模型从cpu移到gpu中，有以下两种方法：
use_cuda = torch.cuda.is_available()
# 方法一：
if use_cuda:
data = data.cuda()
model.cuda()
# 方法二：
device = torch.device("cuda" if use_cuda else "cpu")
data = data.to(device)
model.to(device)个人比较习惯第二种方法，可以少一个 if 语句。而且该方法还可以通过设备号指定使用哪个GPU设备，比如使用0号设备：
device = torch.device("cuda:0" if use_cuda else "cpu")
示例3：
G:\deep-learn\deep-learning\pytorch_classification\bilibili-video-main\Pytorch多卡加速训练\作业\mnist_train_dp.py
代码：
import time
import argparse
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
from torch.optim.lr_scheduler import StepLR
from torch.nn.parallel import DataParallel
torch.manual_seed(0)
torch.cuda.manual_seed_all(0)
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv1 = nn.Conv2d(1, 32, 3, 1)
self.conv2 = nn.Conv2d(32, 64, 3, 1)
self.dropout = nn.Dropout(0.5)
self.fc1 = nn.Linear(9216, 128)
self.fc2 = nn.Linear(128, 10)
def forward(self, x):
x = self.conv1(x)
x = F.relu(x)
x = self.conv2(x)
x = F.relu(x)
x = F.max_pool2d(x, 2)
x = torch.flatten(x, 1)
x = self.dropout(x)
x = self.fc1(x)
x = F.relu(x)
x = self.fc2(x)
x = F.log_softmax(x)
return x
def train(args, model, device, train_loader, optimizer, epoch):
model.train()
for idx, (images, targets) in enumerate(train_loader):
images, targets = images.to(device), targets.to(device)
pred = model(images)
loss = F.cross_entropy(pred, targets)
optimizer.zero_grad()
loss.backward()
optimizer.step()
if idx % args.log_interval == 0:
print("Train Time:{}, epoch: {}, step: {}, loss: {}".format(time.strftime("%Y-%m-%d%H:%M:%S"), epoch + 1, idx, loss.item()))
def test(args, model, device, test_loader):
model.eval()
test_loss = 0
test_acc = 0
with torch.no_grad():
for (images, targets) in test_loader:
images, targets = images.to(device), targets.to(device)
pred = model(images)
loss = F.cross_entropy(pred, targets, reduction="sum")
test_loss += loss.item()
pred_label = torch.argmax(pred, dim=1, keepdims=True)
test_acc += pred_label.eq(targets.view_as(pred_label)).sum().item()
test_loss /= len(test_loader.dataset)
test_acc /= len(test_loader.dataset)
print("Test Time:{}, loss: {}, acc: {}".format(time.strftime("%Y-%m-%d%H:%M:%S"), test_loss, test_acc))
def main():
parser = argparse.ArgumentParser(description="MNIST TRAINING")
parser.add_argument('--device_ids', type=str, default='0', help="Training Devices, example: '0,1,2'")
parser.add_argument('--epochs', type=int, default=10, help="Training Epoch")
parser.add_argument('--log_interval', type=int, default=100, help="Log Interval")
args = parser.parse_args()
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307), (0.3081))])
dataset_train = datasets.MNIST('../data', train=True, transform=transform)
dataset_test = datasets.MNIST('../data', train=False, transform=transform)
train_loader = DataLoader(dataset_train, batch_size=8, shuffle=True, num_workers=8)
test_loader = DataLoader(dataset_test, batch_size=8, shuffle=False, num_workers=8)
device_ids = list(map(int, args.device_ids.split(',')))
device = torch.device('cuda:{}'.format(device_ids[0]))
model = Net().to(device)
model = DataParallel(model, device_ids=device_ids, output_device=device)
optimizer = optim.Adam(model.parameters(), lr=1e-4)
# -----------------------------------------------------------#
# DataParallel
# parser.add_argument('--device_ids', type=str, default='0',
# help="Training Devices, example: '0,1,2'")
# device_ids=list(map(int,args.device_ids.split(',')))
# device=torch.device('cuda:{]'.format(device_ids[0]))
# model=Net().to(device)
# model=DataParallel(model,device_ids=device_ids,output_device=device)
# -----------------------------------------------------------#
scheduler = StepLR(optimizer, step_size=1)
for epoch in range(args.epochs):
train(args, model, device, train_loader, optimizer, epoch)
test(args, model, device, test_loader)
scheduler.step()
torch.save(model.state_dict(), 'train.pt')
if __name__ == '__main__':
main()
DDP;
# -----------------------------------------------------------#
# DDP
# import torch.distributed as dist
# from torch.nn.parallel import DistributedDataParallel
# from torch.utils.data.distributed import DistributedSampler
# parser.add_argument('--local_rank',type=int,default=-1)
# device_ids=list(map(int,args.device_ids.split))#3,4,5,6-->rank[0,1,2,3]
# dist.init_process_group(backend='ncll')
# device=torch.device('cuda;{}'.format(args.device_ids))
# model=Net().to(device)
# model=DistributedDataParallel(model,device_ids=[device_ids[args.local_rank]],output_device=device_ids[args.local_rank])
# sampler_train=DistributedSampler(dataset_train)
# train_loader=DataLoader(dataset_train,batchsize=8,num_workers=8,sampler=sampler_train)
# test_loader = DataLoader(dataset_test, batch_size=8, shuffle=False, num_workers=8)
# for epoch in range(args.epochs):
# sampler_train.set_epoch(epoch)
# train(args,model,device,train_loader,optimizer,epoch)
# if args.local_rank==0:
# test(args,model,device,test_loader)
# scheduler.step()
# if args.local_rank==0:
# torch.save(model.state_dict(),'train.pth')
# -----------------------------------------------------------#相关代码：
import time
import argparse
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
from torch.optim.lr_scheduler import StepLR
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel
from torch.utils.data.distributed import DistributedSampler
torch.manual_seed(0)
torch.cuda.manual_seed_all(0)
# -----------------------------------------------------------#
# DDP
# import torch.distributed as dist
# from torch.nn.parallel import DistributedDataParallel
# from torch.utils.data.distributed import DistributedSampler
# parser.add_argument('--local_rank',type=int,default=-1)
# device_ids=list(map(int,args.device_ids.split))#3,4,5,6-->rank[0,1,2,3]
# dist.init_process_group(backend='ncll')
# device=torch.device('cuda;{}'.format(args.device_ids))
# model=Net().to(device)
# model=DistributedDataParallel(model,device_ids=[device_ids[args.local_rank]],output_device=device_ids[args.local_rank])
# sampler_train=DistributedSampler(dataset_train)
# train_loader=DataLoader(dataset_train,batchsize=8,num_workers=8,sampler=sampler_train)
# test_loader = DataLoader(dataset_test, batch_size=8, shuffle=False, num_workers=8)
# for epoch in range(args.epochs):
# sampler_train.set_epoch(epoch)
# train(args,model,device,train_loader,optimizer,epoch)
# if args.local_rank==0:
# test(args,model,device,test_loader)
# scheduler.step()
# if args.local_rank==0:
# torch.save(model.state_dict(),'train.pth')
# -----------------------------------------------------------#
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv1 = nn.Conv2d(1, 32, 3, 1)
self.conv2 = nn.Conv2d(32, 64, 3, 1)
self.dropout = nn.Dropout(0.5)
self.fc1 = nn.Linear(9216, 128)
self.fc2 = nn.Linear(128, 10)
self.fc3 = nn.Linear(128, 128)
def forward(self, x):
x = self.conv1(x)
x = F.relu(x)
x = self.conv2(x)
x = F.relu(x)
x = F.max_pool2d(x, 2)
x = torch.flatten(x, 1)
x = self.dropout(x)
x = self.fc1(x)
x = F.relu(x)
x = self.fc2(x)
x = F.log_softmax(x)
return x
def train(args, model, device, train_loader, optimizer, epoch):
model.train()
for idx, (images, targets) in enumerate(train_loader):
images, targets = images.to(device), targets.to(device)
pred = model(images)
loss = F.cross_entropy(pred, targets)
optimizer.zero_grad()
loss.backward()
optimizer.step()
if idx % args.log_interval == 0 and args.local_rank == 0:
print("Train Time:{}, epoch: {}, step: {}, loss: {}".format(time.strftime("%Y-%m-%d%H:%M:%S"), epoch + 1, idx, loss.item()))
def test(args, model, device, test_loader):
model.eval()
test_loss = 0
test_acc = 0
with torch.no_grad():
for (images, targets) in test_loader:
images, targets = images.to(device), targets.to(device)
pred = model(images)
loss = F.cross_entropy(pred, targets, reduction="sum")
test_loss += loss.item()
pred_label = torch.argmax(pred, dim=1, keepdims=True)
test_acc += pred_label.eq(targets.view_as(pred_label)).sum().item()
test_loss /= len(test_loader.dataset)
test_acc /= len(test_loader.dataset)
print("Test Time:{}, loss: {}, acc: {}".format(time.strftime("%Y-%m-%d%H:%M:%S"), test_loss, test_acc))
def main():
parser = argparse.ArgumentParser(description="MNIST TRAINING")
parser.add_argument('--device_ids', type=str, default='0', help="Training Devices")
parser.add_argument('--epochs', type=int, default=10, help="Training Epoch")
parser.add_argument('--log_interval', type=int, default=100, help="Log Interval")
parser.add_argument('--local_rank', type=int, default=-1, help="DDP parameter, do not modify")
args = parser.parse_args()
device_ids = list(map(int, args.device_ids.split(',')))
dist.init_process_group(backend='nccl')
device = torch.device('cuda:{}'.format(device_ids[args.local_rank]))
torch.cuda.set_device(device)
model = Net().to(device)
model = DistributedDataParallel(model, device_ids=[device_ids[args.local_rank]], output_device=device_ids[args.local_rank], find_unused_parameters=True)
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307), (0.3081))])
dataset_train = datasets.MNIST('../data', train=True, transform=transform)
dataset_test = datasets.MNIST('../data', train=False, transform=transform)
sampler_train = DistributedSampler(dataset_train)
train_loader = DataLoader(dataset_train, batch_size=8, num_workers=8, sampler=sampler_train)
test_loader = DataLoader(dataset_test, batch_size=8, shuffle=False, num_workers=8)
optimizer = optim.Adam(model.parameters(), lr=1e-4)
scheduler = StepLR(optimizer, step_size=1)
for epoch in range(args.epochs):
sampler_train.set_epoch(epoch)
train(args, model, device, train_loader, optimizer, epoch)
if args.local_rank == 0:
test(args, model, device, test_loader)
scheduler.step()
if args.local_rank == 0:
torch.save(model.state_dict(), 'train.pt')
if __name__ == '__main__':
main()
扫一扫
2万+
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
