PyTorch 最佳实践和代码模板
2020-01-07 本文已影响0人
顾北向南
训练模型的推荐代码结构
# import statements
import torch
import torch.nn as nn
from torch.utils import data
from prefetch_generator import BackgroundGenerator
...
# set flags / seeds
torch.backends.cudnn.benchmark = True
np.random.seed(1)
torch.manual_seed(1)
torch.cuda.manual_seed(1)
...
# Start with main code
if __name__ == '__main__':
# argparse for additional flags for experiment
parser = argparse.ArgumentParser(description="Train a network for ...")
...
opt = parser.parse_args()
# add code for datasets (we always use train and validation/ test set)
data_transforms = transforms.Compose([
transforms.Resize((opt.img_size, opt.img_size)),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
train_dataset = datasets.ImageFolder(
root=os.path.join(opt.path_to_data, "train"),
transform=data_transforms)
train_data_loader = data.DataLoader(train_dataset, ...)
test_dataset = datasets.ImageFolder(
root=os.path.join(opt.path_to_data, "test"),
transform=data_transforms)
test_data_loader = data.DataLoader(test_dataset ...)
...
# instantiate network (which has been imported from *networks.py*)
net = MyNetwork(...)
...
# create losses (criterion in pytorch)
criterion_L1 = torch.nn.L1Loss()
...
# if running on GPU and we want to use cuda move model there
use_cuda = torch.cuda.is_available()
if use_cuda:
net = net.cuda()
...
# create optimizers
optim = torch.optim.Adam(net.parameters(), lr=opt.lr)
...
# load checkpoint if needed/ wanted
start_n_iter = 0
start_epoch = 0
if opt.resume:
ckpt = load_checkpoint(opt.path_to_checkpoint) # custom method for loading last checkpoint
net.load_state_dict(ckpt['net'])
start_epoch = ckpt['epoch']
start_n_iter = ckpt['n_iter']
optim.load_state_dict(ckpt['optim'])
print("last checkpoint restored")
...
# if we want to run experiment on multiple GPUs we move the models there
net = torch.nn.DataParallel(net)
...
# typically we use tensorboardX to keep track of experiments
writer = SummaryWriter(...)
# now we start the main loop
n_iter = start_n_iter
for epoch in range(start_epoch, opt.epochs):
# set models to train mode
net.train()
...
# use prefetch_generator and tqdm for iterating through data
pbar = tqdm(enumerate(BackgroundGenerator(train_data_loader, ...)),
total=len(train_data_loader))
start_time = time.time()
# for loop going through dataset
for i, data in pbar:
# data preparation
img, label = data
if use_cuda:
img = img.cuda()
label = label.cuda()
...
# It's very good practice to keep track of preparation time and computation time using tqdm to find any issues in your dataloader
prepare_time = start_time-time.time()
# forward and backward pass
optim.zero_grad()
...
loss.backward()
optim.step()
...
# udpate tensorboardX
writer.add_scalar(..., n_iter)
...
# compute computation time and *compute_efficiency*
process_time = start_time-time.time()-prepare_time
pbar.set_description("Compute efficiency: {:.2f}, epoch: {}/{}:".format(
process_time/(process_time+prepare_time), epoch, opt.epochs))
start_time = time.time()
# maybe do a test pass every x epochs
if epoch % x == x-1:
# bring models to evaluation mode
net.eval()
...
#do some tests
pbar = tqdm(enumerate(BackgroundGenerator(test_data_loader, ...)),
total=len(test_data_loader))
for i, data in pbar:
...
# save checkpoint if needed
...
如何让实验可复现
- 我们建议在代码开头设置以下种子：
np.random.seed(1)
torch.manual_seed(1)
torch.cuda.manual_seed(1)
如何进一步提升训练和推理速度？
- 在Nvidia GPUs上，你可以在代码的开头添加以下行。这将允许cuda后端在第一次执行时优化你的图。但是，要注意，如果改变网络输入/输出张量的大小，那么每次发生变化时，图都会被优化。这可能导致运行非常慢和内存不足错误。只有当输入和输出总是相同的形状时才设置此标志。通常情况下，这将导致大约20%的改善。
torch.backends.cudnn.benchmark = True
使用tqdm + prefetch_generator模式计算效率的最佳值是什么？
- 这取决于使用的机器、预处理管道和网络大小。在一个1080Ti GPU上使用SSD硬盘，我们看到一个几乎为1.0的计算效率，这是一个理想的场景。如果使用浅(小)网络或慢速硬盘，这个数字可能会下降到0.1-0.2左右，这取决于你的设置。
即使我没有足够的内存，我如何让batch size > 1？
- 在PyTorch中，我们可以很容易地实现虚拟batch sizes。我们只是不让优化器每次都更新参数，并把batch_size个梯度加起来。
...
# in the main loop
out = net(input)
loss = criterion(out, label)
# we just call backward to sum up gradients but don't perform step here
loss.backward()
total_loss += loss.item() / batch_size
if n_iter % batch_size == batch_size-1:
# here we perform out optimization step using a virtual batch size
optim.step()
optim.zero_grad()
print('Total loss: ', total_loss)
total_loss = 0.0
...
在训练过程中如何调整学习率？
- 我们可以直接使用实例化的优化器得到学习率，如下所示：
...
for param_group in optim.param_groups:
old_lr = param_group['lr']
new_lr = old_lr * 0.1
param_group['lr'] = new_lr
print('Updated lr from {} to {}'.format(old_lr, new_lr))
...
在训练中如何使用一个预训练的模型作为损失(没有后向传播)
- 如果你想使用一个预先训练好的模型，如VGG来计算损失，但不训练它(例如在style-transfer/GANs/Auto-encoder中的感知损失)，你可以使用以下模式
...
# instantiate the model
pretrained_VGG = VGG19(...)
# disable gradients (prevent training)
for p in pretrained_VGG.parameters(): # reset requires_grad
p.requires_grad = False
...
# you don't have to use the no_grad() namespace but can just run the model
# no gradients will be computed for the VGG model
out_real = pretrained_VGG(input_a)
out_fake = pretrained_VGG(input_b)
loss = any_criterion(out_real, out_fake)
...
在PyTorch找那个为什么要用.train()和 .eval()？
- 这些方法用于将BatchNorm2d或Dropout2d等层从训练模式设置为推理模式。每个模块都继承自nn.Module有一个名为istrain的属性。.eval()和.train()只是简单地将这个属性设置为True/ False。有关此方法如何实现的详细信息，请参阅PyTorch中的module代码。
我的模型在推理过程中使用了大量内存/如何在PyTorch中正确运行推理模型？
- 确保在代码执行期间没有计算和存储梯度。你可以简单地使用以下模式来确保：
with torch.no_grad():
# run model here
out_tensor = net(in_tensor)
如何微调预训练模型？
- 在PyTorch你可以冻结层。这将防止在优化步骤中更新它们。
# you can freeze whole modules using
for p in pretrained_VGG.parameters(): # reset requires_grad
p.requires_grad = False
PyTorch在C++上比Python快吗？
- C++版本的速度快10%
PyTorch代码使用cudnn.benchmark=True会变快吗？
- 根据我们的经验，你可以获得约20%的加速。但是，第一次运行模型需要相当长的时间来构建优化的图。在某些情况下(前向传递中的循环、没有固定的输入形状、前向中的if/else等等)，这个标志可能会导致内存不足或其他错误。
PyTorch中的.detach()是怎么工作的？
- 如果从计算图中释放一个张量，这里有一个很好的图解：http://www.bnikolic.co.uk/blog/pytorch-detach.html
