├── net
    ├── __init__.py
    └── model.py
├── .gitignore
├── resultOfCycleGAN.png
├── resultOfUNet-VGG.png
├── data
    ├── src
    │   ├── 1004_966_512_2_1_0_0_.png
    │   ├── 1007_821_512_2_1_0_0_.png
    │   ├── 1008_784_549_2_2_0_0_.png
    │   ├── 1008_784_549_2_2_0_1_.png
    │   ├── 1008_784_549_2_2_1_0_.png
    │   ├── 1008_784_549_2_2_1_1_.png
    │   ├── 1009_798_540_2_2_0_0_.png
    │   └── 1009_798_540_2_2_0_1_.png
    ├── target
    │   ├── 1004_966_512_2_1_0_0_.png
    │   ├── 1007_821_512_2_1_0_0_.png
    │   ├── 1008_784_549_2_2_0_0_.png
    │   ├── 1008_784_549_2_2_0_1_.png
    │   ├── 1008_784_549_2_2_1_0_.png
    │   ├── 1008_784_549_2_2_1_1_.png
    │   ├── 1009_798_540_2_2_0_0_.png
    │   └── 1009_798_540_2_2_0_1_.png
    ├── renameCutedImageOfSrcAndTarget.py
    ├── verifySizeInSrcAndTargetWhereNameIsSame.py
    └── cutImageToEqualSize.py
├── README.md
├── utils
    └── dataset.py
└── train.py


/net/__init__.py:
--------------------------------------------------------------------------------
1 | from .model import VGG16UNet
2 | 


--------------------------------------------------------------------------------
/.gitignore:
--------------------------------------------------------------------------------
1 | *.pyc
2 | __pycache__/
3 | net/__pycache__/
4 | utils/__pycache__/


--------------------------------------------------------------------------------
/resultOfCycleGAN.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/gvfjyy/imageDenoise-UNet-VGG/HEAD/resultOfCycleGAN.png


--------------------------------------------------------------------------------
/resultOfUNet-VGG.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/gvfjyy/imageDenoise-UNet-VGG/HEAD/resultOfUNet-VGG.png


--------------------------------------------------------------------------------
/data/src/1004_966_512_2_1_0_0_.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/gvfjyy/imageDenoise-UNet-VGG/HEAD/data/src/1004_966_512_2_1_0_0_.png


--------------------------------------------------------------------------------
/data/src/1007_821_512_2_1_0_0_.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/gvfjyy/imageDenoise-UNet-VGG/HEAD/data/src/1007_821_512_2_1_0_0_.png


--------------------------------------------------------------------------------
/data/src/1008_784_549_2_2_0_0_.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/gvfjyy/imageDenoise-UNet-VGG/HEAD/data/src/1008_784_549_2_2_0_0_.png


--------------------------------------------------------------------------------
/data/src/1008_784_549_2_2_0_1_.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/gvfjyy/imageDenoise-UNet-VGG/HEAD/data/src/1008_784_549_2_2_0_1_.png


--------------------------------------------------------------------------------
/data/src/1008_784_549_2_2_1_0_.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/gvfjyy/imageDenoise-UNet-VGG/HEAD/data/src/1008_784_549_2_2_1_0_.png


--------------------------------------------------------------------------------
/data/src/1008_784_549_2_2_1_1_.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/gvfjyy/imageDenoise-UNet-VGG/HEAD/data/src/1008_784_549_2_2_1_1_.png


--------------------------------------------------------------------------------
/data/src/1009_798_540_2_2_0_0_.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/gvfjyy/imageDenoise-UNet-VGG/HEAD/data/src/1009_798_540_2_2_0_0_.png


--------------------------------------------------------------------------------
/data/src/1009_798_540_2_2_0_1_.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/gvfjyy/imageDenoise-UNet-VGG/HEAD/data/src/1009_798_540_2_2_0_1_.png


--------------------------------------------------------------------------------
/data/target/1004_966_512_2_1_0_0_.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/gvfjyy/imageDenoise-UNet-VGG/HEAD/data/target/1004_966_512_2_1_0_0_.png


--------------------------------------------------------------------------------
/data/target/1007_821_512_2_1_0_0_.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/gvfjyy/imageDenoise-UNet-VGG/HEAD/data/target/1007_821_512_2_1_0_0_.png


--------------------------------------------------------------------------------
/data/target/1008_784_549_2_2_0_0_.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/gvfjyy/imageDenoise-UNet-VGG/HEAD/data/target/1008_784_549_2_2_0_0_.png


--------------------------------------------------------------------------------
/data/target/1008_784_549_2_2_0_1_.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/gvfjyy/imageDenoise-UNet-VGG/HEAD/data/target/1008_784_549_2_2_0_1_.png


--------------------------------------------------------------------------------
/data/target/1008_784_549_2_2_1_0_.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/gvfjyy/imageDenoise-UNet-VGG/HEAD/data/target/1008_784_549_2_2_1_0_.png


--------------------------------------------------------------------------------
/data/target/1008_784_549_2_2_1_1_.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/gvfjyy/imageDenoise-UNet-VGG/HEAD/data/target/1008_784_549_2_2_1_1_.png


--------------------------------------------------------------------------------
/data/target/1009_798_540_2_2_0_0_.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/gvfjyy/imageDenoise-UNet-VGG/HEAD/data/target/1009_798_540_2_2_0_0_.png


--------------------------------------------------------------------------------
/data/target/1009_798_540_2_2_0_1_.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/gvfjyy/imageDenoise-UNet-VGG/HEAD/data/target/1009_798_540_2_2_0_1_.png


--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
 1 | # imageDenoiseUsingUNet-VGG
 2 | 拓片图像去噪网络，以UNet为基本框架，编码器基于VGG16
 3 | 
 4 | 
 5 | ## 环境：
 6 | Python 3.5.3 <br>
 7 | Pytorch 1.2.0 <br>
 8 | 
 9 | 
10 | ## 使用方法：
11 | 1、准备成对的数据集，以相同的名字分别放置在data/src和data/target文件夹下 <br>
12 | 2、如果图像的尺寸过大&&电脑内存不充足，可利用data/文件夹下的脚本进行剪切处理 <br>
13 | 3、使用train.py进行训练 <br>
14 | 
15 | 
16 | ## 实验结果：
17 | 数据集：甲骨拓片数据集。 <br>
18 | 对比网络：CycleGAN <br>
19 | UNet-VGG的效果优于CycleGAN，可以看出CycleGAN对于字体有弱化现象，且对于细节的处理没有UNet-VGG好。 <br>
20 | ### CycleGAN处理效果
21 | ![image](https://github.com/libai-github/imageDenoiseUsingUNet-VGG/blob/master/resultOfCycleGAN.png)
22 | ### UNet-VGG处理效果
23 | ![image](https://github.com/libai-github/imageDenoiseUsingUNet-VGG/blob/master/resultOfUNet-VGG.png)
24 | 
25 | 


--------------------------------------------------------------------------------
/data/renameCutedImageOfSrcAndTarget.py:
--------------------------------------------------------------------------------
 1 | import os
 2 | from PIL import Image
 3 | 
 4 | if __name__ == "__main__":        
 5 |     a='src'
 6 |     b='target'
 7 |     id=1000
 8 |     for i in os.listdir(a):
 9 |         img=Image.open(os.path.join(a,i))
10 |         img.convert('L')
11 |         os.remove(os.path.join(a,i))
12 |         img.save(a+'/'+str(id)+'.png')
13 |         
14 |         img=Image.open(os.path.join(b,i))
15 |         img.convert('L')
16 |         os.remove(os.path.join(b,i))
17 |         img.save(b+'/'+str(id)+'.png')
18 |         
19 |         
20 |         print(id)
21 |         id+=1
22 |         
23 |         
24 |         
25 |         
26 |         
27 | 
28 | 
29 | 
30 | 
31 | 
32 | 


--------------------------------------------------------------------------------
/data/verifySizeInSrcAndTargetWhereNameIsSame.py:
--------------------------------------------------------------------------------
 1 | import os
 2 | from PIL import Image
 3 | import numpy as np
 4 | import shutil
 5 | 
 6 | if __name__ == "__main__":
 7 |     
 8 |     a='src'
 9 |     b='target'
10 |     
11 |     for i in os.listdir(a):
12 |         imga=Image.open(os.path.join(a,i))
13 |         imgb=Image.open(os.path.join(b,i))
14 |         if imga.size!=imgb.size:
15 |             print(imga.size,imgb.size)
16 |             #imgb=imgb.resize(imga.size)
17 |             #imgb.save(os.path.join(b,i))
18 |             #os.remove(os.path.join(a,i))
19 |             #os.remove(os.path.join(b,i))
20 |             #shutil.move(os.path.join(a,i), os.path.join(c,i))
21 |             #shutil.move(os.path.join(b,i), os.path.join(d,i))
22 |             print(i)
23 |     '''
24 |     
25 |     a='src'
26 |     b='target'
27 |     num=0
28 |     for i in os.listdir(a):
29 |         
30 |         imga=Image.open(os.path.join(a,i))
31 |         imgb=Image.open(os.path.join(b,i))
32 |         if imga.size[0]<300 or imga.size[1]<300:
33 |             print(imga.size,imgb.size)
34 |             print(i)
35 |             num+=1
36 |     print(num)
37 |         
38 |     '''
39 |     
40 |     
41 |         
42 | 
43 | 
44 | 
45 | 
46 | 
47 | 
48 | 
49 | 


--------------------------------------------------------------------------------
/data/cutImageToEqualSize.py:
--------------------------------------------------------------------------------
 1 | import os
 2 | from PIL import Image
 3 | import numpy as np
 4 | import math
 5 | 
 6 | if __name__ == "__main__":
 7 |     root='src'
 8 |     size=512
 9 |     for i in os.listdir(root):
10 |         path=os.path.join(root,i)
11 |         img=Image.open(path).convert('L')
12 |         if img.size[0]<512:
13 |             a,b=img.size
14 |             img=img.resize((512,int(b*(512/a))))
15 |             print('chuxian')
16 |         if img.size[1]<512:
17 |             a,b=img.size
18 |             img=img.resize((int(a*(512/b)),512))
19 |             print('chuxian')
20 |         img=np.array(img)
21 |         print(img.shape)
22 |         
23 |         h,w=img.shape
24 |         y_max=math.ceil(h/size)
25 |         x_max=math.ceil(w/size)
26 |         for y in range(y_max):
27 |             for x in range(x_max):
28 |                 if (y+1==y_max) and (x+1<x_max):
29 |                     Image.fromarray(img[h-size:h,x*size:(x+1)*size]).save(path[0:len(path)-4]+'_'+str(h)+'_'+str(w)+'_'+str(y_max)+'_'+str(x_max)+'_'+str(y)+'_'+str(x)+'_.png')
30 |                 elif (x+1==x_max) and (y+1<y_max):
31 |                     Image.fromarray(img[y*size:(y+1)*size,w-size:w]).save(path[0:len(path)-4]+'_'+str(h)+'_'+str(w)+'_'+str(y_max)+'_'+str(x_max)+'_'+str(y)+'_'+str(x)+'_.png')
32 |                 elif (y+1==y_max) and (x+1==x_max):
33 |                     Image.fromarray(img[h-size:h,w-size:w]).save(path[0:len(path)-4]+'_'+str(h)+'_'+str(w)+'_'+str(y_max)+'_'+str(x_max)+'_'+str(y)+'_'+str(x)+'_.png')
34 |                 else:
35 |                     Image.fromarray(img[y*size:(y+1)*size,x*size:(x+1)*size]).save(path[0:len(path)-4]+'_'+str(h)+'_'+str(w)+'_'+str(y_max)+'_'+str(x_max)+'_'+str(y)+'_'+str(x)+'_.png')
36 |         os.remove(path)
37 | 
38 |         
39 | 
40 | 
41 | 
42 | 
43 | 
44 | 
45 | 
46 | 


--------------------------------------------------------------------------------
/utils/dataset.py:
--------------------------------------------------------------------------------
 1 | from os.path import splitext
 2 | import numpy as np
 3 | from glob import glob
 4 | import torch
 5 | from torch.utils.data import Dataset
 6 | import logging
 7 | import os
 8 | import random
 9 | from PIL import Image
10 | 
11 | class MyDataset(Dataset):
12 |     def __init__(self, dir_src ,dir_target ,out_channel, scale):
13 |         self.dir_src=dir_src
14 |         self.dir_target=dir_target
15 |         self.scale = scale
16 |         self.out_channel=out_channel
17 |         self.full=os.listdir(dir_src)
18 |         self.img_src=[]
19 |         self.img_target=[]
20 |         for i in self.full:
21 |             self.img_src.append(os.path.join(dir_src,i))
22 |             self.img_target.append(os.path.join(dir_target,i))
23 |             
24 |     def __len__(self):
25 |         return len(self.full)
26 |     
27 |     def __getitem__(self, i):
28 |         src=Image.open(self.img_src[i]).convert('RGB')
29 |         target=Image.open(self.img_target[i]).convert('RGB')
30 |         if self.out_channel==1:
31 |             src=src.convert('L')
32 |             target=target.convert('L')
33 |         src,target= self.preprocess(src, target, self.out_channel, self.scale)
34 |         #转化为tensor
35 |         return {'src': torch.from_numpy(src), 'target': torch.from_numpy(target)}
36 |     
37 |     def preprocess(cls, img1 ,img2 , out_channel, scale):
38 |         w,h=img1.size
39 |         if random.randint(0,2)==0:
40 |             t=random.randint(0,359)
41 |             img1=img1.rotate(t)#旋转
42 |             img1=img1.crop((int(w*0.15), int(h*0.15), int(w*0.85), int(h*0.85)))#裁剪
43 |             img1=img1.resize((w,h))#恢复原大小  
44 |             img2=img2.rotate(t)#旋转
45 |             img2=img2.crop((int(w*0.15), int(h*0.15), int(w*0.85), int(h*0.85)))#裁剪
46 |             img2=img2.resize((w,h))#恢复原大小  
47 | 
48 |         t=random.randint(0,10)#翻转
49 |         if t==0:
50 |             img1=img1.transpose(Image.FLIP_LEFT_RIGHT)
51 |             img2=img2.transpose(Image.FLIP_LEFT_RIGHT)
52 |         elif t==1:
53 |             img1=img1.transpose(Image.FLIP_TOP_BOTTOM)
54 |             img2=img2.transpose(Image.FLIP_TOP_BOTTOM)
55 |         elif t==2:
56 |             img1=img1.transpose(Image.ROTATE_90)
57 |             img2=img2.transpose(Image.ROTATE_90)
58 |         elif t==3:
59 |             img1=img1.transpose(Image.ROTATE_180)
60 |             img2=img2.transpose(Image.ROTATE_180)
61 |         elif t==4:
62 |             img1=img1.transpose(Image.ROTATE_270)
63 |             img2=img2.transpose(Image.ROTATE_270)
64 |         elif t==5:
65 |             img1=img1.transpose(Image.TRANSPOSE)
66 |             img2=img2.transpose(Image.TRANSPOSE)
67 |         elif t==6:
68 |             img1=img1.transpose(Image.TRANSVERSE)
69 |             img2=img2.transpose(Image.TRANSVERSE)
70 |         
71 |         img1=np.array(img1)
72 |         img2=np.array(img2)
73 |         if out_channel==3:
74 |             # HWC to CHW
75 |             img1 = np.array(img1).transpose((2, 0, 1))
76 |             img2 = np.array(img2).transpose((2, 0, 1))
77 |         elif out_channel==1:
78 |             img1=img1.reshape(-1,img1.shape[0],img1.shape[1])
79 |             img2=img2.reshape(-1,img2.shape[0],img2.shape[1])
80 | 
81 |         #归一化至[0,1]        
82 |         img1 = (img1 / 255)
83 |         img2 = (img2 / 255)
84 |         return img1, img2
85 | 


--------------------------------------------------------------------------------
/net/model.py:
--------------------------------------------------------------------------------
  1 | from torch import nn
  2 | import torch
  3 | from torchvision import models
  4 | import torchvision
  5 | from torch.nn import functional as F
  6 | 
  7 | 
  8 | def conv3x3(in_, out):
  9 |     return nn.Conv2d(in_, out, 3, padding=1)
 10 | 
 11 | 
 12 | class ConvRelu(nn.Module):
 13 |     def __init__(self, in_: int, out: int):
 14 |         super(ConvRelu, self).__init__()
 15 |         self.conv = conv3x3(in_, out)
 16 |         self.activation = nn.ReLU(inplace=True)
 17 | 
 18 |     def forward(self, x):
 19 |         x = self.conv(x)
 20 |         x = self.activation(x)
 21 |         return x
 22 | 
 23 | 
 24 | class DecoderBlock(nn.Module):
 25 |     """
 26 |     Paramaters for Deconvolution were chosen to avoid artifacts, following
 27 |     link https://distill.pub/2016/deconv-checkerboard/
 28 |     """
 29 | 
 30 |     def __init__(self, in_channels, middle_channels, out_channels, is_deconv=True):
 31 |         super(DecoderBlock, self).__init__()
 32 |         self.in_channels = in_channels
 33 | 
 34 |         if is_deconv:
 35 |             self.block = nn.Sequential(
 36 |                 ConvRelu(in_channels, middle_channels),
 37 |                 nn.ConvTranspose2d(middle_channels, out_channels, kernel_size=4, stride=2,
 38 |                                    padding=1),
 39 |                 nn.ReLU(inplace=True)
 40 |             )
 41 |         else:
 42 |             self.block = nn.Sequential(
 43 |                 nn.Upsample(scale_factor=2, mode='bilinear'),
 44 |                 ConvRelu(in_channels, middle_channels),
 45 |                 ConvRelu(middle_channels, out_channels),
 46 |             )
 47 | 
 48 |     def forward(self, x):
 49 |         return self.block(x)
 50 | 
 51 | 
 52 | class VGG16UNet(nn.Module):
 53 |     def __init__(self, out_channel=1, num_filters=32, pretrained=False):
 54 |         """
 55 |         out_channel:输出图像的通道数
 56 |         pretrained：是否加载预训练模型
 57 |         """
 58 |         super().__init__()
 59 |         self.out_channel = out_channel
 60 | 
 61 |         self.pool = nn.MaxPool2d(2, 2)
 62 | 
 63 |         self.encoder = torchvision.models.vgg16(pretrained=pretrained).features#预训练模型会自动下载
 64 | 
 65 |         self.relu = nn.ReLU(inplace=True)
 66 | 
 67 |         self.encoder_conv1_1=nn.Conv2d(out_channel, 64, 3, padding=(1, 1))#修改VGG16的首层以使用单通道图像，此部分参数。
 68 |         '''
 69 |         self.conv1 = nn.Sequential(self.encoder_conv1_1,
 70 |                                    self.relu,
 71 |                                    self.encoder[2],
 72 |                                    self.relu)
 73 |         '''
 74 |         self.conv1 = nn.Sequential(self.encoder[0],
 75 |                                    self.relu,
 76 |                                    self.encoder[2],
 77 |                                    self.relu)
 78 | 
 79 |         self.conv2 = nn.Sequential(self.encoder[5],
 80 |                                    self.relu,
 81 |                                    self.encoder[7],
 82 |                                    self.relu)
 83 | 
 84 |         self.conv3 = nn.Sequential(self.encoder[10],
 85 |                                    self.relu,
 86 |                                    self.encoder[12],
 87 |                                    self.relu,
 88 |                                    self.encoder[14],
 89 |                                    self.relu)
 90 | 
 91 |         self.conv4 = nn.Sequential(self.encoder[17],
 92 |                                    self.relu,
 93 |                                    self.encoder[19],
 94 |                                    self.relu,
 95 |                                    self.encoder[21],
 96 |                                    self.relu)
 97 | 
 98 |         self.conv5 = nn.Sequential(self.encoder[24],
 99 |                                    self.relu,
100 |                                    self.encoder[26],
101 |                                    self.relu,
102 |                                    self.encoder[28],
103 |                                    self.relu)
104 | 
105 |         self.center = DecoderBlock(512, num_filters * 8 * 2, num_filters * 8)
106 | 
107 |         self.dec5 = DecoderBlock(512 + num_filters * 8, num_filters * 8 * 2, num_filters * 8)
108 |         self.dec4 = DecoderBlock(512 + num_filters * 8, num_filters * 8 * 2, num_filters * 8)
109 |         self.dec3 = DecoderBlock(256 + num_filters * 8, num_filters * 4 * 2, num_filters * 2)
110 |         self.dec2 = DecoderBlock(128 + num_filters * 2, num_filters * 2 * 2, num_filters)
111 |         self.dec1 = ConvRelu(64 + num_filters, num_filters)
112 |         self.final = nn.Conv2d(num_filters, out_channel, kernel_size=1)
113 | 
114 |     def forward(self, x):
115 |         conv1 = self.conv1(x)
116 |         conv2 = self.conv2(self.pool(conv1))
117 |         conv3 = self.conv3(self.pool(conv2))
118 |         conv4 = self.conv4(self.pool(conv3))
119 |         conv5 = self.conv5(self.pool(conv4))
120 | 
121 |         center = self.center(self.pool(conv5))
122 | 
123 |         dec5 = self.dec5(torch.cat([center, conv5], 1))
124 | 
125 |         dec4 = self.dec4(torch.cat([dec5, conv4], 1))
126 |         dec3 = self.dec3(torch.cat([dec4, conv3], 1))
127 |         dec2 = self.dec2(torch.cat([dec3, conv2], 1))
128 |         dec1 = self.dec1(torch.cat([dec2, conv1], 1))
129 |         
130 |         x_out=F.relu(self.final(dec1))
131 |         #x_out=self.final(dec1)
132 | 
133 |         return x_out
134 |         
135 | 


--------------------------------------------------------------------------------
/train.py:
--------------------------------------------------------------------------------
  1 | import argparse
  2 | import os
  3 | import sys
  4 | 
  5 | import numpy as np
  6 | import torch
  7 | import torch.nn as nn
  8 | from torch import optim
  9 | 
 10 | from net import VGG16UNet
 11 | from loss import My_loss
 12 | 
 13 | from utils.dataset import MyDataset
 14 | from torch.utils.data import DataLoader
 15 | 
 16 | from torchsummary import summary
 17 | 
 18 | np.set_printoptions(threshold = np.inf) 
 19 | np.set_printoptions(suppress = True)
 20 | from PIL import Image  
 21 | 
 22 | 
 23 | dir_src = 'data/src/'
 24 | dir_target = 'data/target/'
 25 | dir_checkpoint = 'checkpoints/'
 26 | 
 27 | 
 28 | def train_net(net,
 29 |               device,
 30 |               epochs=5,
 31 |               batch_size=1,
 32 |               lr=0.001,
 33 |               save_cp=True,
 34 |               img_scale=0.5,
 35 |               out_channel=3):
 36 |     #数据加载器
 37 |     dataset = MyDataset(dir_src, dir_target, out_channel, img_scale)
 38 |     train_loader = DataLoader(dataset, batch_size=batch_size, shuffle=True, num_workers=2, pin_memory=True)
 39 |     #优化器损失函数
 40 |     optimizer = optim.Adam(net.parameters(),lr=lr)
 41 |     criterion = torch.nn.MSELoss()
 42 | 
 43 |     
 44 |     for epoch in range(101,epochs):
 45 |         #调整学习率
 46 |         if epoch % 20 == 0:
 47 |             for p in optimizer.param_groups:
 48 |                 p['lr'] *= 0.5
 49 | 
 50 | 
 51 |         net.train()
 52 |         epoch_loss = 0
 53 |         global_step=0
 54 |         for batch in train_loader:
 55 |             #获取数据并加载到CPU/GPU
 56 |             src = batch['src']
 57 |             target = batch['target']
 58 |             src = src.to(device=device, dtype=torch.float32)
 59 |             target = target.to(device=device, dtype=torch.float32)  
 60 | 
 61 |             optimizer.zero_grad()   
 62 |             src_pred = net(src)#前向传播结果
 63 |             '''
 64 |             dis=torch.tensor(src_pred).cpu().numpy()
 65 |             print('src_pred:',np.sum(dis))
 66 |             print(dis[0][0][0])
 67 |             #print(dis[0][0][0])
 68 |             
 69 |             
 70 |             if global_step==5:
 71 |                 image_src=(torch.tensor(src).cpu().numpy()[0][0])*255
 72 |                 image_target=(torch.tensor(target).cpu().numpy()[0][0])*255
 73 |                 img_src=Image.fromarray(image_src)
 74 |                 img_src.show()
 75 |                 img_target=Image.fromarray(image_target)
 76 |                 img_target.show()
 77 |             '''
 78 |             
 79 |             
 80 |             loss = criterion(src_pred, target)#计算损失
 81 |             epoch_loss += loss.item()
 82 |             loss.backward()
 83 |             optimizer.step()
 84 |             
 85 |             if global_step%100==0:
 86 |                 print(src_pred[0][0][0][0])
 87 |                 print('global_step:',global_step,' ; loss:',loss.item())
 88 |             global_step+=1
 89 |             
 90 |         print('-------------------------------------------epoch:',epoch,' ; loss_mean:',epoch_loss/len(train_loader))
 91 |  
 92 |         #保存模型     
 93 |         if save_cp and epoch%10==0:
 94 |             try:
 95 |                 os.mkdir(dir_checkpoint)
 96 |                 print('Created checkpoint directory')
 97 |             except OSError:
 98 |                 pass
 99 |             torch.save(net.state_dict(),
100 |                        dir_checkpoint + f'Checkpoint_epoch{epoch}.pth')
101 |             print(f'Checkpoint {epoch} saved !')
102 | 
103 | 
104 | def get_args():
105 |     parser = argparse.ArgumentParser(description='Train the VGG16UNet on srcimages and target',
106 |                                      formatter_class=argparse.ArgumentDefaultsHelpFormatter)
107 |     parser.add_argument('-e', '--epochs', metavar='E', type=int, default=500,
108 |                         help='Number of epochs', dest='epochs')
109 |     parser.add_argument('-b', '--batch-size', metavar='B', type=int, nargs='?', default=1,
110 |                         help='Batch size', dest='batchsize')
111 |     parser.add_argument('-l', '--learning-rate', metavar='LR', type=float, nargs='?', default=0.0001,
112 |                         help='Learning rate', dest='lr')
113 |     parser.add_argument('-f', '--load', dest='load', type=str, default=True,
114 |                         help='Load model from a .pth file')
115 |     parser.add_argument('-p', '--pretrained',dest='pretrained', type=str, default=False,
116 |                         help='Load pretrained model from internet')
117 |     parser.add_argument('-s', '--scale', dest='scale', type=float, default=0.5,
118 |                         help='Downscaling factor of the images')
119 |     parser.add_argument('-o', '--out_channel', dest='out_channel', type=int, default=3,
120 |                         help='channel of the images')
121 |     return parser.parse_args()
122 | 
123 | 
124 | if __name__ == '__main__':
125 |     args = get_args()
126 |     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
127 |     #可自动加载预训练模型的参数：将事先下载好的模型直接复制到缓存目录以节省临时下载耗费时间。
128 |     net = VGG16UNet(out_channel=args.out_channel, pretrained=args.pretrained)
129 |     
130 |     #初始化Conv2d,ReLU,MaxPool2d,ConvTranspose2d,
131 |     num=0
132 |     for m in net.modules():
133 |         if isinstance(m,nn.Conv2d) and num>14:#在此之前为VGG16参数，不用再初始化
134 |             torch.nn.init.xavier_normal_(m.weight.data)#Xavier 正态分布
135 |             if m.bias is not None:
136 |                 m.bias.data.zero_()
137 |         elif  isinstance(m,nn.ReLU):
138 |             continue
139 |         elif isinstance(m,nn.MaxPool2d):
140 |             continue
141 |         elif isinstance(m,nn.ConvTranspose2d):
142 |             torch.nn.init.xavier_normal_(m.weight.data)#Xavier 正态分布
143 |             if m.bias is not None:
144 |                 m.bias.data.zero_()
145 |         num+=1
146 | 
147 |     #加载上次训练的参数
148 |     net.load_state_dict(torch.load('Checkpoint_epoch100l1.pth'))
149 |     
150 | 
151 |     #加载数据到CPU/GPU
152 |     net.to(device=device,dtype=torch.float32)
153 |     #summary(net, (1, 224, 224)) 
154 | 
155 |     try:
156 |         train_net(net=net,
157 |                   epochs=args.epochs,
158 |                   batch_size=args.batchsize,
159 |                   lr=args.lr,
160 |                   device=device,
161 |                   img_scale=args.scale,
162 |                   out_channel=args.out_channel
163 |                   )
164 |     except KeyboardInterrupt:
165 |         torch.save(net.state_dict(), 'INTERRUPTED.pth')
166 |         print('Saved interrupt')
167 |         try:
168 |             sys.exit(0)
169 |         except SystemExit:
170 |             os._exit(0)
171 | 


--------------------------------------------------------------------------------