PyTorch Cookbook（常用代码段整理合集）
https://zhuanlan.zhihu.com/p/59205847?
1.freezn_bn
使用imagenet预训练模型时，可以对BN层的参数进行固定，这样再别的小数据集上也可以得到比较好的结果，同时还可以加快训练速度。
def set_mode(self, mode, is_freeze_bn=False ):
self.mode = mode
if mode in ['eval', 'valid', 'test']:
self.eval()
elif mode in ['backup']:
self.train()
if is_freeze_bn==True: ##freeze
for m in self.modules():
if isinstance(m, BatchNorm2d):
m.eval()
m.weight.requires_grad = False
m.bias.requires_grad = False
2.输入归一化方式
可以通过数据集计算均值和方差，也可以通过对数据进行BN操作
if self.is_first_bn:
x = self.first_bn(x)
else:
mean=[0.485, 0.456, 0.406] #rgb
std =[0.229, 0.224, 0.225]
x = torch.cat([
(x[:,[0]]-mean[0])/std[0],
(x[:,[1]]-mean[1])/std[1],
(x[:,[2]]-mean[2])/std[2],
],1)
3.将预训练模型CPU格式转成GPU格式
一般预训练模型保存的格式都是CPU下的，当需要加载到GPU上训练时，需要在每一层的名字上加上module.，否则会报不匹配错误。
def load_pretrain(self, pretrain_file):
pretrain_state_dict = torch.load(pretrain_file)
state_dict = self.state_dict()
keys = list(state_dict.keys())
for key in keys:
state_dict[key] = pretrain_state_dict['module.'+key]
self.load_state_dict(state_dict)
print('load: '+pretrain_file)
4. cos循环学习率
import torch
import math
from torch.optim.lr_scheduler import _LRScheduler
class CosineAnnealingLR_with_Restart(_LRScheduler):
"""Set the learning rate of each parameter group using a cosine annealing
schedule, where :math:`\eta_{max}` is set to the initial lr and
:math:`T_{cur}` is the number of epochs since the last restart in SGDR:
.. math::
\eta_t = \eta_{min} + \frac{1}{2}(\eta_{max} - \eta_{min})(1 +
\cos(\frac{T_{cur}}{T_{max}}\pi))
When last_epoch=-1, sets initial lr as lr.
It has been proposed in
`SGDR: Stochastic Gradient Descent with Warm Restarts`_. The original pytorch
implementation only implements the cosine annealing part of SGDR,
I added my own implementation of the restarts part.
Args:
optimizer (Optimizer): Wrapped optimizer.
T_max (int): Maximum number of iterations.
T_mult (float): Increase T_max by a factor of T_mult
eta_min (float): Minimum learning rate. Default: 0.
last_epoch (int): The index of last epoch. Default: -1.
model (pytorch model): The model to save.
out_dir (str): Directory to save snapshots
take_snapshot (bool): Whether to save snapshots at every restart
.. _SGDR\: Stochastic Gradient Descent with Warm Restarts:
https://arxiv.org/abs/1608.03983
"""
def __init__(self, optimizer, T_max, T_mult, model, out_dir, take_snapshot, eta_min=0, last_epoch=-1):
self.T_max = T_max
self.T_mult = T_mult
self.Te = self.T_max
self.eta_min = eta_min
self.current_epoch = last_epoch
self.model = model
self.out_dir = out_dir
self.take_snapshot = take_snapshot
self.lr_history = []
super(CosineAnnealingLR_with_Restart, self).__init__(optimizer, last_epoch)
def get_lr(self):
new_lrs = [self.eta_min + (base_lr - self.eta_min) *
(1 + math.cos(math.pi * self.current_epoch / self.Te)) / 2
for base_lr in self.base_lrs]
self.lr_history.append(new_lrs)
return new_lrs
def step(self, epoch=None):
if epoch is None:
epoch = self.last_epoch + 1
self.last_epoch = epoch
self.current_epoch += 1
for param_group, lr in zip(self.optimizer.param_groups, self.get_lr()):
param_group['lr'] = lr
## restart
if self.current_epoch == self.Te:
print("restart at epoch {:03d}".format(self.last_epoch + 1))
if self.take_snapshot:
torch.save({
'epoch': self.T_max,
'state_dict': self.model.state_dict()
}, self.out_dir + "Weight/" + 'snapshot_e_{:03d}.pth.tar'.format(self.T_max))
## reset epochs since the last reset
self.current_epoch = 0
## reset the next goal
self.Te = int(self.Te * self.T_mult)
self.T_max = self.T_max + self.Te
使用方法
optimizer = optim.SGD(filter(lambda p: p.requires_grad, net.parameters()),
lr=0.1, momentum=0.9, weight_decay=0.0005)
sgdr = CosineAnnealingLR_with_Restart(optimizer,
T_max=config.cycle_inter,
T_mult=1,
model=net,
out_dir='../input/',
take_snapshot=False,
eta_min=1e-3)
for cycle_index in range(config.cycle_num):
print('cycle index: ' + str(cycle_index))
min_acer = 1.0
5.类别不平衡问题
balance: 对正例和负例分别计算loss，使得正例loss权重之和与负例loss权重之和相等，解决数据不平衡的问题。
随机每次多从list中随机取
if self.mode == 'train':
if self.balance:
if random.randint(0,1)==0:
tmp_list = self.train_list[0]
else:
tmp_list = self.train_list[1]
pos = random.randint(0,len(tmp_list)-1)
color, depth, ir, label = tmp_list[pos]
else:
color, depth, ir, label = self.train_list[index]
6.将三种图片进行concat叠加变成6通道
image = np.concatenate([color.reshape([n, self.image_size, self.image_size, 3]),
depth.reshape([n, self.image_size, self.image_size, 3]),
ir.reshape([n, self.image_size, self.image_size, 3])],
axis=3)
image = np.transpose(image, (0, 3, 1, 2))
image = image.astype(np.float32)
image = image.reshape([n, self.channels * 3, self.image_size, self.image_size])
image = image / 255.0
7.数据增强
https://github.com/SeuTao/CVPR19-Face-Anti-spoofing/blob/master/process/augmentation.py
7.1 随机crop
def random_cropping(image, target_shape=(32, 32, 3), is_random = True):
image = cv2.resize(image,(RESIZE_SIZE,RESIZE_SIZE))
target_h, target_w,_ = target_shape
height, width, _ = image.shape
if is_random:
start_x = random.randint(0, width - target_w)
start_y = random.randint(0, height - target_h)
else:
start_x = ( width - target_w ) // 2
start_y = ( height - target_h ) // 2
zeros = image[start_y:start_y+target_h,start_x:start_x+target_w,:]
return zeros
7.2 随机裁剪并叠加
def TTA_18_cropps(image, target_shape=(32, 32, 3)):
image = cv2.resize(image, (RESIZE_SIZE, RESIZE_SIZE))
width, height, d = image.shape
target_w, target_h, d = target_shape
start_x = ( width - target_w) // 2
start_y = ( height - target_h) // 2
starts = [[start_x, start_y],
[start_x - target_w, start_y],
[start_x, start_y - target_w],
[start_x + target_w, start_y],
[start_x, start_y + target_w],
[start_x + target_w, start_y + target_w],
[start_x - target_w, start_y - target_w],
[start_x - target_w, start_y + target_w],
[start_x + target_w, start_y - target_w],
]
images = []
for start_index in starts:
image_ = image.copy()
x, y = start_index
if x < 0:
x = 0
if y < 0:
y = 0
if x + target_w >= RESIZE_SIZE:
x = RESIZE_SIZE - target_w-1
if y + target_h >= RESIZE_SIZE:
y = RESIZE_SIZE - target_h-1
zeros = image_[x:x + target_w, y: y+target_h, :]
image_ = zeros.copy()
zeros = np.fliplr(zeros)
image_flip = zeros.copy()
images.append(image_.reshape([1,target_shape[0],target_shape[1],target_shape[2]]))
images.append(image_flip.reshape([1,target_shape[0],target_shape[1],target_shape[2]]))
return images
7.4 随机添加块噪声
def random_erasing(img, probability = 0.5, sl = 0.02, sh = 0.5, r1 = 0.5, channel = 3):
if random.uniform(0, 1) > probability:
return img
for attempt in range(100):
area = img.shape[0] * img.shape[1]
target_area = random.uniform(sl, sh) * area
aspect_ratio = random.uniform(r1, 1 / r1)
h = int(round(math.sqrt(target_area * aspect_ratio)))
w = int(round(math.sqrt(target_area / aspect_ratio)))
if w < img.shape[1] and h < img.shape[0]:
x1 = random.randint(0, img.shape[0] - h)
y1 = random.randint(0, img.shape[1] - w)
noise = np.random.random((h,w,channel))*255
noise = noise.astype(np.uint8)
if img.shape[2] == channel:
img[x1:x1 + h, y1:y1 + w, :] = noise
else:
print('wrong')
return
return img
return img
7.5计算模型的flops
https://github.com/Tramac/torchscope
``