1.搭建AlexNet网络
alexnet_model.py
# -*-coding:utf-8-*-
import torch.nn as nn
import torch
class AlexNet(nn.Module):
def __init__(self, num_classes=1000, init_weights=False):
super(AlexNet, self).__init__()
self.features = nn.Sequential(
nn.Conv2d(3, 48, kernel_size=11, stride=4, padding=2), # input[3, 224, 224] output[48, 55, 55]
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=3, stride=2), # output[48, 27, 27]
nn.Conv2d(48, 128, kernel_size=5, padding=2), # output[128, 27, 27]
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=3, stride=2), # output[128, 13, 13]
nn.Conv2d(128, 192, kernel_size=3, padding=1), # output[192, 13, 13]
nn.ReLU(inplace=True),
nn.Conv2d(192, 192, kernel_size=3, padding=1), # output[192, 13, 13]
nn.ReLU(inplace=True),
nn.Conv2d(192, 128, kernel_size=3, padding=1), # output[128, 13, 13]
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=3, stride=2), # output[128, 6, 6]
)
self.classifier = nn.Sequential(
nn.Dropout(p=0.5),
nn.Linear(128 * 6 * 6, 2048),
nn.ReLU(inplace=True),
nn.Dropout(p=0.5),
nn.Linear(2048, 2048),
nn.ReLU(inplace=True),
nn.Linear(2048, num_classes),
)
if init_weights:
self._initialize_weights()
def forward(self, x):
x = self.features(x)
x = torch.flatten(x, start_dim=1)
x = self.classifier(x)
return x
def _initialize_weights(self):
for m in self.modules():
if isinstance(m, nn.Conv2d):
nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
if m.bias is not None:
nn.init.constant_(m.bias, 0)
elif isinstance(m, nn.Linear):
nn.init.normal_(m.weight, 0, 0.01)
nn.init.constant_(m.bias, 0)
2.训练网络
alexnet_train.py
# -*-coding:utf-8-*-
import os
import json
import torch
import torch.nn as nn
from torchvision import transforms, datasets, utils
import torch.optim as optim
from tqdm import tqdm
from alexnet_model import AlexNet
def main():
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print("using {} device.".format(device))
data_transform = {
"train": transforms.Compose([transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]),
"val": transforms.Compose([transforms.Resize((224, 224)), # cannot 224, must (224, 224)
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])}
data_root = os.path.abspath(os.path.join(os.getcwd(), "../")) # get data root path,脚本位置
image_path = os.path.join(data_root, "data_set") # rock data set path
assert os.path.exists(image_path), "{} path does not exist.".format(image_path)
train_dataset = datasets.ImageFolder(root=os.path.join(image_path, "train"),
transform=data_transform["train"])
train_num = len(train_dataset)
rock_list = train_dataset.class_to_idx
cla_dict = dict((val, key) for key, val in rock_list.items())
# write dict into json file
json_str = json.dumps(cla_dict, indent=4)
with open('class_indices.json', 'w') as json_file:
json_file.write(json_str)
batch_size = 32
nw = min([os.cpu_count(), batch_size if batch_size > 1 else 0, 8]) # number of workers
print('Using {} dataloader workers every process'.format(nw))
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=batch_size, shuffle=True,
num_workers=nw)
validate_dataset = datasets.ImageFolder(root=os.path.join(image_path, "val"),
transform=data_transform["val"])
val_num = len(validate_dataset)
validate_loader = torch.utils.data.DataLoader(validate_dataset,
batch_size=4, shuffle=False,
num_workers=nw)
print("using {} images for training, {} images for validation.".format(train_num,
val_num))
net = AlexNet(num_classes=7, init_weights=True)
net.to(device)
loss_function = nn.CrossEntropyLoss()
# pata = list(net.parameters())
optimizer = optim.Adam(net.parameters(), lr=0.0002)
epochs = 10
save_path = './AlexNet.pth'
best_acc = 0.0
train_steps = len(train_loader)
for epoch in range(epochs):
# train
net.train()
running_loss = 0.0
train_bar = tqdm(train_loader)
for step, data in enumerate(train_bar):
images, labels = data
optimizer.zero_grad()
outputs = net(images.to(device))
loss = loss_function(outputs, labels.to(device))
loss.backward()
optimizer.step()
# print statistics
running_loss += loss.item()
train_bar.desc = "train epoch[{}/{}] loss:{:.3f}".format(epoch + 1,
epochs,
loss)
# validate
net.eval()
acc = 0.0 # accumulate accurate number / epoch
with torch.no_grad():
val_bar = tqdm(validate_loader)
for val_data in val_bar:
val_images, val_labels = val_data
outputs = net(val_images.to(device))
predict_y = torch.max(outputs, dim=1)[1]
acc += torch.eq(predict_y, val_labels.to(device)).sum().item()
val_accurate = acc / val_num
print('[epoch %d] train_loss: %.3f val_accuracy: %.3f' %
(epoch + 1, running_loss / train_steps, val_accurate))
if val_accurate > best_acc:
best_acc = val_accurate
torch.save(net.state_dict(), save_path)
print('Finished Training')
if __name__ == '__main__':
main()
3.训练迭代结果
/home/xyjin/anaconda3/envs/pytorch/bin/python /home/xyjin/PycharmProjects/data_mining/data_process/rock/classed_data/alexnet/alexnet_train.py
using cpu device.
Using 4 dataloader workers every process
using 12600 images for training, 1400 images for validation.
train epoch[1/10] loss:1.057: 100%|██████████| 394/394 [06:53<00:00, 1.05s/it]
100%|██████████| 350/350 [00:22<00:00, 15.39it/s]
[epoch 1] train_loss: 1.281 val_accuracy: 0.629
train epoch[2/10] loss:0.865: 100%|██████████| 394/394 [06:52<00:00, 1.05s/it]
100%|██████████| 350/350 [00:23<00:00, 14.86it/s]
[epoch 2] train_loss: 1.015 val_accuracy: 0.678
train epoch[3/10] loss:1.058: 100%|██████████| 394/394 [06:59<00:00, 1.06s/it]
100%|██████████| 350/350 [00:22<00:00, 15.55it/s]
[epoch 3] train_loss: 0.923 val_accuracy: 0.714
train epoch[4/10] loss:0.676: 100%|██████████| 394/394 [06:59<00:00, 1.07s/it]
100%|██████████| 350/350 [00:22<00:00, 15.23it/s]
[epoch 4] train_loss: 0.876 val_accuracy: 0.723
train epoch[5/10] loss:1.047: 100%|██████████| 394/394 [06:58<00:00, 1.06s/it]
100%|██████████| 350/350 [00:22<00:00, 15.31it/s]
[epoch 5] train_loss: 0.817 val_accuracy: 0.745
train epoch[6/10] loss:1.065: 100%|██████████| 394/394 [06:59<00:00, 1.07s/it]
100%|██████████| 350/350 [00:22<00:00, 15.49it/s]
[epoch 6] train_loss: 0.767 val_accuracy: 0.741
train epoch[7/10] loss:0.653: 100%|██████████| 394/394 [07:00<00:00, 1.07s/it]
100%|██████████| 350/350 [00:22<00:00, 15.77it/s]
[epoch 7] train_loss: 0.741 val_accuracy: 0.803
train epoch[8/10] loss:0.668: 100%|██████████| 394/394 [07:00<00:00, 1.07s/it]
100%|██████████| 350/350 [00:22<00:00, 15.69it/s]
[epoch 8] train_loss: 0.674 val_accuracy: 0.831
train epoch[9/10] loss:0.373: 100%|██████████| 394/394 [07:01<00:00, 1.07s/it]
100%|██████████| 350/350 [00:23<00:00, 15.14it/s]
[epoch 9] train_loss: 0.662 val_accuracy: 0.766
train epoch[10/10] loss:0.610: 100%|██████████| 394/394 [06:57<00:00, 1.06s/it]
100%|██████████| 350/350 [00:22<00:00, 15.62it/s]
[epoch 10] train_loss: 0.616 val_accuracy: 0.835
Finished Training
Process finished with exit code 0
4.需要注意的地方
导入数据集相关代码
...
data_root = os.path.abspath(os.path.join(os.getcwd(), "../")) # get data root path,脚本位置
image_path = os.path.join(data_root, "data_set") # rock data set path
assert os.path.exists(image_path), "{} path does not exist.".format(image_path)
train_dataset = datasets.ImageFolder(root=os.path.join(image_path, "train"),
transform=data_transform["train"])
...
1.data_root = os.path.abspath(os.path.join(os.getcwd(), “…/”))
- os.getcwd()表示获取当前脚本所在的位置
- "…/"表示返回上一层目录
- 脚本所在位置：/home/xyjin/anaconda3/envs/pytorch/bin/python /home/xyjin/PycharmProjects/data_mining/data_process/rock/classed_data/alexnet/alexnet_train.py
- 则返回上一层目录就是在alexnet
2.image_path = os.path.join(data_root, “data_set”) # rock data set path
- os.path.join表示把data_root和data_set结合起来
- image_path:/home/xyjin/anaconda3/envs/pytorch/bin/python /home/xyjin/PycharmProjects/data_mining/data_process/rock/classed_data/alexnet/data_set
6303
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
