pytorch REPRODUCIBILITY学习
环境
torch 1.11.0+cu113
torchvision 0.12.0+cu113
numpy 1.21.4
imgaug 0.4.0
Pillow 9.2.0
opencv-python 4.3.0.36
opencv-python-headless 4.2.0.34
REPRODUCIBILITY
随机种子
首先要设置随机种子
#!/usr/bin/env python
# _*_ coding:utf-8 _*_
import os
import random
import numpy as np
import torch
def setup_seed(seed):
"""
set random seed
:param seed: seed num
"""
os.environ['PYTHONHASHSEED'] = str(seed)
os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":16:8" # LSTM(cuda>10.2)
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
torch.cuda.manual_seed_all(seed) # if you are using multi-GPU.
torch.use_deterministic_algorithms(True)
# torch.backends.cudnn.enabled = False
torch.backends.cudnn.benchmark = False
torch.backends.cudnn.deterministic = True
不过如果你的库用的随机种子不是numpy，那需要自己设置
设置了torch.use_deterministic_algorithms(True)之后，如果你的代码里有不确定算法，就会报错
在这里可以看你的算法是不是不确定的https://pytorch.org/docs/stable/generated/torch.use_deterministic_algorithms.html#torch.use_deterministic_algorithms
多个num_workers
然后是dataloader，如果你的num_workers>1,则需要
def seed_worker(worker_id):
worker_seed = torch.initial_seed() % 2**32
numpy.random.seed(worker_seed)
random.seed(worker_seed)
g = torch.Generator()
g.manual_seed(0)
DataLoader(
train_dataset,
batch_size=batch_size,
num_workers=num_workers,
worker_init_fn=seed_worker,
generator=g,
)
下面测试只测试1个num_worker
dataloader顺序
主要测试一下固定种子后，获取图片顺序是否一样
utils/utils.py
用来设置随机种子
#!/usr/bin/env python
# _*_ coding:utf-8 _*_
import os
import random
import numpy as np
import torch
def setup_seed(seed):
"""
set random seed
:param seed: seed num
"""
os.environ['PYTHONHASHSEED'] = str(seed)
os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":16:8" # LSTM(cuda>10.2)
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
torch.cuda.manual_seed_all(seed) # if you are using multi-GPU.
torch.use_deterministic_algorithms(True)
# torch.backends.cudnn.enabled = False
torch.backends.cudnn.benchmark = False
torch.backends.cudnn.deterministic = True
def seed_worker(worker_id):
worker_seed = torch.initial_seed() % 2 ** 32
np.random.seed(worker_seed)
random.seed(worker_seed)
dataset/landmark_dataset.py
#!/usr/bin/env python
# _*_ coding:utf-8 _*_
import os
from glob import glob
import imgaug as ia
import imgaug.augmenters as iaa
import numpy as np
import torch
from cv2 import cv2
from imgaug.augmentables import KeypointsOnImage
from torch.utils.data import Dataset
class LandmarkDataset(Dataset):
def __init__(self, dataset_type='train', perform_augmentation=True):
self.landmark_num = 19
junior_annotations = glob(
os.path.join('/data/datasets/ISBI2015_ceph', 'AnnotationsByMD', '400_junior', '*.txt'))
junior_annotations.sort()
senior_annotations = glob(
os.path.join('/data/datasets/ISBI2015_ceph', 'AnnotationsByMD', '400_senior', '*.txt'))
senior_annotations.sort()
img_paths = glob(
os.path.join('/data/datasets/ISBI2015_ceph', 'RawImage', 'TrainingData', '*.bmp'))
img_paths.sort()
if 'train' == dataset_type:
self.img_paths = img_paths[:130]
self.junior_annotations = junior_annotations[:130]
self.senior_annotations = senior_annotations[:130]
elif 'val' == dataset_type:
self.img_paths = img_paths[130:151]
self.junior_annotations = junior_annotations[130:150]
self.senior_annotations = senior_annotations[130:150]
self.perform_augmentation = perform_augmentation
self.height = 512
self.width = 416
# Define augmentation
seq = [
iaa.Resize({"width": self.width, "height": self.height}),
]
if self.perform_augmentation:
seq.append(iaa.Sometimes(0.9, [
iaa.Affine(translate_px={"x": (-10, 10),
"y": (-10, 10)},
scale=(0.95, 1),
rotate=(-3, 3)),
iaa.Multiply(mul=(0.5, 1.5)),
iaa.GammaContrast(),
iaa.ElasticTransformation(alpha=(0, 500),
sigma=30, order=3)
]))
self.augmentation = iaa.Sequential(seq)
def __len__(self):
return len(self.img_paths)
def __getitem__(self, idx):
print(self.img_paths[idx])
img = cv2.imread(self.img_paths[idx], 0)
junior_annotation = np.loadtxt(self.junior_annotations[idx], delimiter=",", max_rows=self.landmark_num)
senior_annotation = np.loadtxt(self.senior_annotations[idx], delimiter=",", max_rows=self.landmark_num)
# mean landmark
annotation = ((junior_annotation + senior_annotation) / 2).astype(int)
kps_array = KeypointsOnImage.from_xy_array(annotation, img.shape)
# aug
img_aug, landmark_aug = self.augmentation(image=img, keypoints=kps_array)
# draw landmark
img_aug = landmark_aug.draw_on_image(cv2.cvtColor(img_aug, cv2.COLOR_GRAY2RGB), color=(0, 255, 0), size=3)
return torch.FloatTensor(img_aug.transpose((2, 0, 1))).contiguous()/255.0
main.py
#!/usr/bin/env python
# _*_ coding:utf-8 _*_
import argparse
import os
import torch
from torch.utils.data import DataLoader
from torchvision import utils as vutils
from dataset.landmark_dataset import LandmarkDataset
from utils.utils import setup_seed, seed_worker
g = torch.Generator()
g.manual_seed(0)
def get_args():
parser = argparse.ArgumentParser()
# optional
parser.add_argument("-s", "--save_dir")
return parser.parse_args()
if __name__ == '__main__':
args = get_args()
setup_seed(42)
train_dataset = LandmarkDataset()
train_data_loader = DataLoader(train_dataset,
batch_size=4,
shuffle=True,
worker_init_fn=seed_worker,
generator=g)
valid_dataset = LandmarkDataset(dataset_type='val')
valid_data_loader = DataLoader(valid_dataset,
batch_size=4,
shuffle=True,
worker_init_fn=seed_worker,
generator=g)
for epoch in range(10):
for idx, x in enumerate(train_data_loader):
cur_dir = f'{args.save_dir}/{epoch}/{idx}'
os.makedirs(cur_dir, exist_ok=True)
for i, t in enumerate(x):
vutils.save_image(t.unsqueeze(0), os.path.join(cur_dir, f'{i}.bmp'))
print()
if 0 == (epoch + 1) % 5:
for x in valid_data_loader:
pass
文件结构
测试方法
输出两次获取图片的顺序，然后对比
python main.py --save_dir imgaug1>test1
python main.py --save_dir imgaug2>test2
diff test1 test2
也就是说生成顺序是一样的
测试imgaug数据增强
根据imgaug作者在https://github.com/aleju/imgaug/issues/555说的
imgagu会先从numpy获取第一个随机种子，所以imgaug前设置numpy随机种子，也就相当于设置了imgaug的随机种子
check_same.py
#!/usr/bin/env python
# _*_ coding:utf-8 _*_
import argparse
import os
import numpy as np
from cv2 import cv2
def is_same(direction1, direction2):
for root, dirs, files in os.walk(direction1):
for file in files:
path1 = os.path.join(root, file)
path2 = path1.replace(direction1, direction2)
img1 = cv2.imread(path1)
img2 = cv2.imread(path2)
if not np.allclose(img1, img2):
return False
return True
def get_args():
parser = argparse.ArgumentParser()
# optional
parser.add_argument("--d1")
parser.add_argument("--d2")
return parser.parse_args()
if __name__ == '__main__':
args = get_args()
print(is_same(args.d1, args.d2))
其他程序跟上面一样
测试方法
python main.py --save_dir imgaug1
python main.py --save_dir imgaug2
python check_same.py --d1 imgaug1 --d2 imgaug2
也就是说跑2次，增强的是一样的
to_deterministic
如果用了这个，那你只有第一下是随机的，后面的增强都跟第一次一样
不过如果你传入一个batch，那batch之间是可以不同的
比如说第一个batch，旋转5，10，15
那么第二个batch就也是旋转5，10，15
模型
model/unet2d.py
import torch
import torch.nn as nn
import torch.nn.functional as F
class DoubleConv(nn.Module):
"""(convolution => [BN] => ReLU) * 2"""
def __init__(self, in_channels, out_channels, mid_channels=None):
super().__init__()
if not mid_channels:
mid_channels = out_channels
self.double_conv = nn.Sequential(
nn.Conv2d(in_channels, mid_channels, kernel_size=3, padding=1),
nn.BatchNorm2d(mid_channels),
nn.ReLU(inplace=True),
nn.Conv2d(mid_channels, out_channels, kernel_size=3, padding=1),
nn.BatchNorm2d(out_channels),
nn.ReLU(inplace=True)
)
def forward(self, x):
return self.double_conv(x)
class Down(nn.Module):
"""Downscaling with maxpool then double conv"""
def __init__(self, in_channels, out_channels):
super().__init__()
self.maxpool_conv = nn.Sequential(
nn.MaxPool2d(2),
DoubleConv(in_channels, out_channels)
)
def forward(self, x):
return self.maxpool_conv(x)
class Up(nn.Module):
"""Upscaling then double conv"""
def __init__(self, in_channels, out_channels, bilinear=True):
super().__init__()
# if bilinear, use the normal convolutions to reduce the number of channels
if bilinear:
self.up = nn.Upsample(
scale_factor=2, mode='bilinear', align_corners=True)
# self.up = nn.Upsample(scale_factor=2)
self.conv = DoubleConv(in_channels, out_channels, in_channels // 2)
else:
self.up = nn.ConvTranspose2d(
in_channels, in_channels // 2, kernel_size=2, stride=2)
self.conv = DoubleConv(in_channels, out_channels)
def forward(self, x1, x2):
x1 = self.up(x1)
# input is CHW
diffY = x2.size()[2] - x1.size()[2]
diffX = x2.size()[3] - x1.size()[3]
x1 = F.pad(x1, [diffX // 2, diffX - diffX // 2,
diffY // 2, diffY - diffY // 2])
# if you have padding issues, see
# https://github.com/HaiyongJiang/U-Net-Pytorch-Unstructured-Buggy/commit/0e854509c2cea854e247a9c615f175f76fbb2e3a
# https://github.com/xiaopeng-liao/Pytorch-UNet/commit/8ebac70e633bac59fc22bb5195e513d5832fb3bd
x = torch.cat([x2, x1], dim=1)
return self.conv(x)
class OutConv(nn.Module):
def __init__(self, in_channels, out_channels):
super(OutConv, self).__init__()
self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1)
def forward(self, x):
return self.conv(x)
class UNet(nn.Module):
def __init__(self, in_channels, out_channels, bilinear=True):
super(UNet, self).__init__()
if not isinstance(in_channels, list):
in_channels = [in_channels]
if not isinstance(out_channels, list):
out_channels = [out_channels]
self.in_channels = in_channels
self.out_channels = out_channels
self.bilinear = bilinear
for i, (n_chan, n_class) in enumerate(zip(in_channels, out_channels)):
setattr(self, 'in{i}'.format(i=i), OutConv(n_chan, 64))
setattr(self, 'out{i}'.format(i=i), OutConv(64, n_class))
self.conv = DoubleConv(64, 64)
self.down1 = Down(64, 128)
self.down2 = Down(128, 256)
self.down3 = Down(256, 512)
factor = 2 if bilinear else 1
self.down4 = Down(512, 1024 // factor)
self.up1 = Up(1024, 512 // factor, bilinear)
self.up2 = Up(512, 256 // factor, bilinear)
self.up3 = Up(256, 128 // factor, bilinear)
self.up4 = Up(128, 64, bilinear)
def forward(self, x, task_idx=0):
x1 = getattr(self, 'in{}'.format(task_idx))(x)
x1 = self.conv(x1)
x2 = self.down1(x1)
x3 = self.down2(x2)
x4 = self.down3(x3)
x5 = self.down4(x4)
x = self.up1(x5, x4)
x = self.up2(x, x3)
x = self.up3(x, x2)
x = self.up4(x, x1)
logits = getattr(self, 'out{}'.format(task_idx))(x)
return {'output': torch.sigmoid(logits)}
顺带一提，nn.Upsample(scale_factor=2, mode=‘bilinear’, align_corners=True)是在cuda里不确定算法
（可以用下面的代码试试
#!/usr/bin/env python
# _*_ coding:utf-8 _*_
import os
import random
import numpy as np
import torch
from torch import nn
def setup_seed(seed):
"""
set random seed
:param seed: seed num
"""
os.environ['PYTHONHASHSEED'] = str(seed)
os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":16:8" # LSTM(cuda>10.2)
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
torch.cuda.manual_seed_all(seed) # if you are using multi-GPU.
torch.use_deterministic_algorithms(True, warn_only=True)
# torch.backends.cudnn.enabled = False
torch.backends.cudnn.benchmark = False
torch.backends.cudnn.deterministic = True
setup_seed(42)
device = torch.device('cuda')
if __name__ == '__main__':
x = torch.randn((1, 3, 64, 64), requires_grad=True).to(device)
target = torch.randn((1, 3, 128, 128)).to(device)
m = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True).to(device)
y_hat = m(x)
loss = torch.sum((y_hat - target) ** 2)
print(loss)
loss.backward()
改成nn.Upsample(scale_factor=2)的话，是确定算法
save_model.py
#!/usr/bin/env python
# _*_ coding:utf-8 _*_
import argparse
import torch
from model.unet2d import UNet
from utils.utils import setup_seed
def get_args():
parser = argparse.ArgumentParser()
# optional
parser.add_argument("-s", "--save_name")
return parser.parse_args()
if __name__ == '__main__':
args = get_args()
setup_seed(42)
model = UNet(1, 19)
torch.save(model, args.save_name)
check_same_model.py
#!/usr/bin/env python
# _*_ coding:utf-8 _*_
import argparse
import torch
def get_args():
parser = argparse.ArgumentParser()
# optional
parser.add_argument("--m1")
parser.add_argument("--m2")
return parser.parse_args()
def is_same(path1, path2):
model1 = torch.load(path1)
model2 = torch.load(path2)
for p1, p2 in zip(model1.parameters(), model2.parameters()):
if not torch.allclose(p1, p2):
return False
return True
if __name__ == '__main__':
args = get_args()
print(is_same(args.m1, args.m2))
其他跟上面一样
测试方法
简单来说就是生成两个，然后逐个参数对比
python save_model.py --save_name model1.pth
python save_model.py --save_name model2.pth
python check_same_model.py --m1 model1.pth --m2 model2.pth
如果你有不确定算法
可以看到两次结果不一样
如果你都是确定算法
这3次结果都是一样的
参考
https://pytorch.org/docs/stable/notes/randomness.html
https://zhuanlan.zhihu.com/p/480931657
https://github.com/aleju/imgaug/issues/555