1 # 导入所需的包
2 import torch
3 #import wandb
4 import torch.nn as nn
5 from torchvision import transforms
6 from torch.utils.data import DataLoader
7 from torchvision.datasets import CIFAR100
8
9 # 使用Compose容器组合定义图像预处理方式
10 transf = transforms.Compose([
11 # 改变图像大小
12 transforms.Resize(224),
13 # 将给定图片转为shape为(C, H, W)的tensor
14 transforms.ToTensor()
15 ])
16 # 数据准备
17 traindata = CIFAR100(
18 # 数据集的地址
19 root="./",
20 # 是否为训练集，True为训练集
21 train=True,
22 # 使用数据预处理
23 transform=transf,
24 # 是否需要下载， True为需要下载
25 download=True
26 )
27 testdata = CIFAR100(
28 root="./",
29 train=False,
30 transform=transf,
31 download=True
32 )
33 # 定义数据加载器
34 trainloader = DataLoader(
35 # 需要加载的数据
36 traindata,
37 # 定义batch大小
38 batch_size=128,
39 # 是否打乱顺序，True为打乱顺序
40 shuffle=True
41 )
42 testloader = DataLoader(
43 testdata,
44 batch_size=128,
45 shuffle=False
46 )
47
48 # 定义AlexNet网络
49 # 此处写代码
50 from torch import nn
51 class AlexNet(nn.Module):
52 # 初始化
53 def __init__(self):
54 super(AlexNet, self).__init__()
55 self.conv1 = nn.Conv2d(in_channels=3,
56 out_channels=96,
57 kernel_size=11,
58 padding=2,
59 stride=4)
60 self.relu1 = nn.ReLU()
61 self.max_pool1 = nn.MaxPool2d(kernel_size=3,
62 stride=2,
63 padding=0)
64
65 self.conv2 = nn.Conv2d(in_channels=96,
66 out_channels=256,
67 kernel_size=5,
68 padding=2,
69 stride=1)
70 self.relu2 = nn.ReLU()
71 self.max_pool2 = nn.MaxPool2d(kernel_size=3,
72 stride=2,
73 padding=0)
74
75 self.conv3 = nn.Conv2d(in_channels=256,
76 out_channels=384,
77 kernel_size=3,
78 padding=1,
79 stride=1)
80 self.relu3 = nn.ReLU()
81
82 self.conv4 = nn.Conv2d(in_channels=384,
83 out_channels=384,
84 kernel_size=3,
85 padding=1,
86 stride=1)
87 self.relu4 = nn.ReLU()
88
89 self.conv5 = nn.Conv2d(in_channels=384,
90 out_channels=256,
91 kernel_size=3,
92 padding=1,
93 stride=1)
94 self.relu5 = nn.ReLU()
95 self.max_pool5 = nn.MaxPool2d(kernel_size=3,
96 stride=2,
97 padding=0)
98
99 self.dropout1=nn.Dropout(0.5)
100 self.linear1 = nn.Linear(in_features=256*6*6,
101 out_features=4096,
102 bias=True)
103 self.relu6 = nn.ReLU()
104
105 self.dropout2 = nn.Dropout(0.5)
106 self.linear2 = nn.Linear(in_features=4096,
107 out_features=4096,
108 bias=True)
109 self.relu7 = nn.ReLU()
110
111 self.linear3 = nn.Linear(in_features=4096,
112 out_features=100,
113 bias=True)
114
115
116 # 定义前向计算过程
117 def forward(self, x):
118 x = self.conv1(x)
119 x = self.relu1(x)
120 x = self.max_pool1(x)
121
122 x = self.conv2(x)
123 x = self.relu2(x)
124 x = self.max_pool2(x)
125
126 x = self.conv3(x)
127 x = self.relu3(x)
128
129 x = self.conv4(x)
130 x = self.relu4(x)
131
132 x = self.conv5(x)
133 x = self.relu5(x)
134 x = self.max_pool5(x)
135
136 # 将特征展平（超级重要！！！）
137 x = x.view(x.shape[0], -1)
138
139 x = self.dropout1(x)
140 x = self.linear1(x)
141 x = self.relu6(x)
142
143 x = self.dropout2(x)
144 x = self.linear2(x)
145 x = self.relu7(x)
146
147
148 x = self.linear3(x)
149
150 return x
151
152 # 定义网络的预训练
153 def train(net, train_loader, test_loader, device, l_r = 0.0002, num_epochs=25,):
154 # 使用wandb跟踪训练过程
155 #experiment = wandb.init(project='AlexNet', resume='allow', anonymous='must')
156 # 定义损失函数
157 criterion = nn.CrossEntropyLoss()
158 # 定义优化器
159 optimizer = torch.optim.Adam(net.parameters(), lr=l_r)
160 # 将网络移动到指定设备
161 net = net.to(device)
162 # 正式开始训练
163 for epoch in range(num_epochs):
164 # 保存一个Epoch的损失
165 train_loss = 0
166 # 计算准确度
167 test_corrects = 0
168 # 设置模型为训练模式
169 net.train()
170 for step, (imgs, labels) in enumerate(train_loader):
171 # 训练使用的数据移动到指定设备
172 imgs = imgs.to(device)
173 labels = labels.to(device)
174 output = net(imgs)
175 # 计算损失
176 loss = criterion(output, labels)
177 # 将梯度清零
178 optimizer.zero_grad()
179 # 将损失进行后向传播
180 loss.backward()
181 # 更新网络参数
182 optimizer.step()
183 train_loss += loss.item()
184 # 设置模型为验证模式
185 net.eval()
186 for step, (imgs, labels) in enumerate(test_loader):
187 imgs = imgs.to(device)
188 labels = labels.to(device)
189 output = net(imgs)
190 pre_lab = torch.argmax(output, 1)
191 corrects = (torch.sum(pre_lab == labels.data).double() / imgs.size(0))
192 test_corrects += corrects.item()
193 #一个Epoch结束时，使用wandb保存需要可视化的数据
194 # experiment.log({
195 # 'epoch':epoch,
196 # 'train loss': train_loss / len(train_loader),
197 # 'test acc': test_corrects / len(test_loader),
198 # })
199 print('Epoch: {}/{}'.format(epoch, num_epochs-1))
200 print('{} Train Loss:{:.4f}'.format(epoch, train_loss / len(train_loader)))
201 print('{} Test Acc:{:.4f}'.format(epoch, test_corrects / len(test_loader)))
202 # 保存此Epoch训练的网络的参数
203 torch.save(net.state_dict(), './net.pth')
204
205 if __name__ == "__main__":
206 # 定义训练使用的设备
207 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
208 net = AlexNet()
209 train(net, trainloader, testloader, device, l_r=0.0003, num_epochs=10)
标签：nn,代码,train,CIFAR100,net,True,self,AlexNet,size From： https://www.cnblogs.com/jevonChao/p/17266546.html