1 # 导入所需的包
2 import torch
3 #import wandb
4 import torch.nn as nn
5 from torchvision import transforms
6 from torchvision.datasets import CIFAR100
7 from torch.utils.data import DataLoader
8
9 # 使用Compose容器组合定义图像预处理方式
10
11
12 transf = transforms.Compose([
13 # 将给定图片转为shape为(C, H, W)的tensor
14 transforms.ToTensor()
15 ])
16 # 数据准备
17 train_set = CIFAR100(
18 # 数据集的地址
19 root="./",
20 # 是否为训练集，True为训练集
21 train=True,
22 # 使用数据预处理
23 transform=transf,
24 # 是否需要下载， True为需要下载
25 download=True
26 )
27 test_set = CIFAR100(
28 root="./",
29 train=False,
30 transform=transf,
31 download=True
32 )
33 # 定义数据加载器
34 train_loader = DataLoader(
35 # 需要加载的数据
36 train_set,
37 # 定义batch大小
38 batch_size=16,
39 # 是否打乱顺序，True为打乱顺序
40 shuffle=True
41 )
42 test_loader = DataLoader(
43 test_set,
44 batch_size=16,
45 shuffle=False
46 )
47
48
49 # 定义VGG16网络
50 # 提示：nn.Sequetial（）：按顺序包装
51 # nn.ModuleList（）：像list一样包装，用[ ]
52 # nn.ModuleDict（）：像dict一样包装，用{ }
53 # 以上几个不能同时使用，可以选其中一种
54 import torch.nn as nn
55 class My_VGG(nn.Module):
56 def __init__(self, num_classes):
57 super(My_VGG, self).__init__()
58 # 第一个卷积块
59 self.conv1 = nn.Sequential(
60 nn.Conv2d(3, 64, kernel_size=(3, 3), stride=1, padding=1),nn.ReLU(),
61 # nn.BatchNorm2d(64)归一化处理，注意括号中是指通道数，
62 # 在样本过少的时候采用GroupNorm或LayerNorm
63 nn.Conv2d(64, 64, kernel_size=(3, 3), stride=1, padding=1),nn.ReLU(),
64 )
65 self.pool1 = nn.MaxPool2d((2, 2), 2)
66 # 第二个卷积块
67 self.conv2 = nn.Sequential(
68 nn.Conv2d(64, 128, kernel_size=(3, 3), stride=1, padding=1),nn.ReLU(),
69 nn.Conv2d(128, 128, kernel_size=(3, 3), stride=1, padding=1),nn.ReLU(),
70 )
71 self.pool2 = nn.MaxPool2d((2, 2), 2)
72 # 第三个卷积块
73 self.conv3 = nn.Sequential(
74 nn.Conv2d(128, 256, kernel_size=(3, 3), stride=1, padding=1),nn.ReLU(),
75 nn.Conv2d(256, 256, kernel_size=(3, 3), stride=1, padding=1),nn.ReLU(),
76 nn.Conv2d(256, 256, kernel_size=(3, 3), stride=1, padding=1),nn.ReLU(),
77 )
78 self.pool3 = nn.MaxPool2d((2, 2), 2)
79 # 第四个卷积块
80 self.conv4 = nn.Sequential(
81 nn.Conv2d(256, 512, kernel_size=(3, 3), stride=1, padding=1),nn.ReLU(),
82 nn.Conv2d(512, 512, kernel_size=(3, 3), stride=1, padding=1),nn.ReLU(),
83 nn.Conv2d(512, 512, kernel_size=(3, 3), stride=1, padding=1),nn.ReLU(),
84 )
85 self.pool4 = nn.MaxPool2d((2, 2), 2)
86 # 第五个卷积块
87 self.conv5 = nn.Sequential(
88 nn.Conv2d(512, 512, kernel_size=(3, 3), stride=1, padding=1),nn.ReLU(),
89 nn.Conv2d(512, 512, kernel_size=(3, 3), stride=1, padding=1),nn.ReLU(),
90 nn.Conv2d(512, 512, kernel_size=(3, 3), stride=1, padding=1),nn.ReLU(),
91 )
92 self.pool5 = nn.MaxPool2d((2, 2), 2)
93 # 全连接层部分
94 self.output = nn.Sequential(
95 nn.Linear(512*1*1, 4096),nn.ReLU(),nn.Dropout(),
96 nn.Linear(4096, 4096),nn.ReLU(),nn.Dropout(0.5),
97 nn.Linear(4096, num_classes)
98 )
99 '''在此处初始化，
100 定义在网络里
101 如下两种：
102 '''
103 #Xavier初始化
104 for m in self.modules():
105 if isinstance(m, nn.Linear) or isinstance(m, nn.Conv2d):
106 torch.nn.init.xavier_uniform_(m.weight)
107 #
108 #He初始化（Kaiming初始化）
109 # for m in self.modules():
110 # if isinstance(m, nn.Linear) or isinstance(m, nn.Conv2d):
111 # torch.nn.init.kaiming_normal_(m.weight)
112
113
114 def forward(self,x):
115 x = self.pool1(self.conv1(x))
116 x = self.pool2(self.conv2(x))
117 x = self.pool3(self.conv3(x))
118 x = self.pool4(self.conv4(x))
119 x = self.pool5(self.conv5(x))
120 x = x.view(x.size(0), -1)
121 outer = self.output(x)
122 return outer
123
124
125 # 定义网络的预训练
126 def train(net, train_loader, test_loader, device, l_r = 0.01, num_epochs=25):
127 # 使用wandb跟踪训练过程
128 #experiment = wandb.init(project='VGG16', resume='allow', anonymous='must')
129 # 定义损失函数
130 criterion = nn.CrossEntropyLoss()
131 # 定义优化器
132 optimizer = torch.optim.Adam(net.parameters(), lr=l_r)
133
134 # 将网络移动到指定设备
135 net = net.to(device)
136 # 正式开始训练
137 for epoch in range(num_epochs):
138 # 保存一个Epoch的损失
139 train_loss = 0
140 # 计算准确度
141 test_corrects = 0
142 # 设置模型为训练模式
143 net.train()
144 for step, (imgs, labels) in enumerate(train_loader):
145 # 训练使用的数据移动到指定设备
146 imgs = imgs.to(device)
147 labels = labels.to(device)
148 output = net(imgs)
149 # 计算损失
150 loss = criterion(output, labels)
151 # 将梯度清零
152 optimizer.zero_grad()
153 # 将损失进行后向传播
154 loss.backward()
155 # 更新网络参数
156 optimizer.step()
157 train_loss += loss.item()
158 pre_lab = torch.argmax(output, 1)
159 train_batch_corrects = (torch.sum(pre_lab == labels.data).double() / imgs.size(0))
160 if step % 100 == 0:
161 print("train {} {}/{} loss: {} acc: {}".format(epoch, step, len(train_loader), loss.item(),
162 train_batch_corrects.item()))
163 # 设置模型为验证模式
164 net.eval()
165 for step, (imgs, labels) in enumerate(test_loader):
166 imgs = imgs.to(device)
167 labels = labels.to(device)
168 output = net(imgs)
169 loss = criterion(output, labels)
170 pre_lab = torch.argmax(output, 1)
171 test_batch_corrects = (torch.sum(pre_lab == labels.data).double() / imgs.size(0))
172 test_corrects += test_batch_corrects.item()
173 if step % 100 == 0:
174 print("val {} {}/{} loss: {} acc: {}".format(epoch, step, len(test_loader), loss.item(),
175 test_batch_corrects.item()))
176
177 # 一个Epoch结束时，使用wandb保存需要可视化的数据
178 # experiment.log({
179 # 'epoch':epoch,
180 # 'train loss': train_loss / len(train_loader),
181 # 'test acc': test_corrects / len(test_loader),
182 # })
183 print('Epoch: {}/{}'.format(epoch, num_epochs-1))
184 print('{} Train Loss:{:.4f}'.format(epoch, train_loss / len(train_loader)))
185 print('{} Test Acc:{:.4f}'.format(epoch, test_corrects / len(test_loader)))
186 # 保存此Epoch训练的网络的参数
187 torch.save(net.state_dict(), './VGG16.pth')
188
189 if __name__ == "__main__":
190 # 定义训练使用的设备
191 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
192 # 使用自定义的VGG16类实现VGG16网络
193 # 由于CIFAR100有100种类别，所以修改VGG网络的输出层神经元数量，num_classes = 100
194
195 My_VGG16=My_VGG
196 net = My_VGG16(num_classes = 100)
197 train(net, train_loader,
198 test_loader, device,
199 #更改学习率
200 l_r=0.00003,
201 num_epochs=50)
标签：nn,VGG16,代码,ReLU,test,train,CIFAR100,self,size From： https://www.cnblogs.com/jevonChao/p/17283481.html