1、权重与模型不匹配时，加载部分模型
# load_state_dict() 默认strict=True，需要完全匹配，否则报错
# 修改为strict=False后，只匹配存在的参数
pretrained_dict = torch.load(weight_path)
model.load_state_dict(pretrained_dict, strict=False)
2、模型中前半部分网络不参与训练
def __init__(self, num_classes=10):
super(self.__class__, self).__init__()
#下面是不参与训练的部分
self.conv1 = nn.Conv2d(3, 64, 3, 1, 1, bias=False)
for p in self.parameters():
p.requires_grad = False
3、在做恢复任务时，有重叠分块处理数据并将处理好的各个块进行组合。
def sliding_forward(net, x, crop_size=800):
n, c, h, w = x.size()
if h <= crop_size and w <= crop_size:
return net(x)[1]
else:
result = torch.zeros(n, c, h, w).cuda()
count = torch.zeros(n, 1, h, w).cuda()
stride = int(crop_size / 3.)
h_steps = 1 + int(ceil(float(max(h - crop_size, 0)) / stride))
w_steps = 1 + int(ceil(float(max(w - crop_size, 0)) / stride))
for h_idx in range(h_steps):
for w_idx in range(w_steps):
ws0, ws1 = w_idx * stride, crop_size + w_idx * stride
hs0, hs1 = h_idx * stride, crop_size + h_idx * stride
if h_idx == h_steps - 1:
hs0, hs1 = max(h - crop_size, 0), h
if w_idx == w_steps - 1:
ws0, ws1 = max(w - crop_size, 0), w
out = net(x[:, :, hs0: hs1, ws0: ws1])[1]
result[:, :, hs0: hs1, ws0: ws1] += out.data
count[:, :, hs0: hs1, ws0: ws1] += 1
assert torch.min(count) > 0
result = result / count
return result
4、数据增强。
#data augmentation for image rotate
def augment(hazy, clean):
augmentation_method = random.choice([0, 1, 2, 3, 4, 5])
rotate_degree = random.choice([90, 180, 270])
'''Rotate'''
if augmentation_method == 0:
hazy = transforms.functional.rotate(hazy, rotate_degree)
clean = transforms.functional.rotate(clean, rotate_degree)
return hazy, clean
'''Vertical'''
if augmentation_method == 1:
vertical_flip = torchvision.transforms.RandomVerticalFlip(p=1)
hazy = vertical_flip(hazy)
clean = vertical_flip(clean)
return hazy, clean
'''Horizontal'''
if augmentation_method == 2:
horizontal_flip = torchvision.transforms.RandomHorizontalFlip(p=1)
hazy = horizontal_flip(hazy)
clean = horizontal_flip(clean)
return hazy, clean
'''no change'''
if augmentation_method == 3 or augmentation_method == 4 or augmentation_method == 5:
return hazy, clean
