1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
| from torchvision import datasets, models, transforms
import os
from torch.autograd import Variable
import torch.utils.data
import random
import matplotlib.pyplot as plt
from PIL import Image
import numpy as np
data_dir = "E:/trains/CatsVSDogs"
data_trainsforms = {
"train": transforms.Compose([
transforms.RandomResizedCrop(300),
transforms.Resize((224, 224)),
transforms.RandomCrop((224, 224)),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]),
]),
"test": transforms.Compose([
transforms.RandomResizedCrop(300),
transforms.RandomHorizontalFlip(),
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]),
]), }
image_datasets = {
x: datasets.ImageFolder(root=os.path.join(data_dir, x),
transform=data_trainsforms[x])
for x in ["train", "test"]
}
data_loader = {x: torch.utils.data.DataLoader(image_datasets[x], batch_size=20, shuffle=True) for x in
["train", "test"]}
X_example, y_example = next(iter(data_loader["train"]))
example_classees = image_datasets["train"].classes
index_classes = image_datasets["train"].class_to_idx
model = models.resnet50(pretrained=True)
Use_gpu = torch.cuda.is_available()
for parma in model.parameters():
parma.requires_grad = False
model.fc = torch.nn.Linear(2048, 2)
if Use_gpu:
model = model.cuda()
loss_f = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.fc.parameters(), lr=0.00001)
epoch_n = 5
for epoch in range(epoch_n):
print("Epoch {}/{}".format(epoch + 1, epoch_n))
print("-" * 10)
for phase in ["train", "test"]:
if phase == "train":
print("训练中。。。")
model.train(True)
else:
print("测试中。。。")
model.train(False)
running_loss = 0.0
running_corrects = 0
for batch, data in enumerate(data_loader[phase], 1):
X, y = data
if Use_gpu:
X, y = Variable(X.cuda()), Variable(y.cuda())
else:
X, y = Variable(X), Variable(y)
y_pred = model(X)
_, pred = torch.max(y_pred.data, 1)
optimizer.zero_grad()
loss = loss_f(y_pred, y)
if phase == "train":
loss.backward()
optimizer.step()
running_loss += loss.item()
running_corrects += torch.sum(pred == y.data)
if batch % 500 == 0 and phase == "train":
print("Batch{},训练损失率:{:.4f},训练正确率:{:.4f}".format(batch, running_loss / batch,
100 * running_corrects / (20 * batch)))
epoch_loss = running_loss * 20 / len(image_datasets[phase])
epoch_acc = 100 * running_corrects / len(image_datasets[phase])
print("{} 当前损失率:{:.4f} 当前正确率:{:.4f}%".format(phase, epoch_loss, epoch_acc))
torch.save(model.state_dict(),'model.ckpt1')
torch.save(model.state_dict(),'model.pth')
print("over")
model.load_state_dict(torch.load('model.pth'))
model.eval()
test_images = image_datasets['test'].imgs
random.shuffle(test_images)
test_images = test_images[:9]
fig, axs = plt.subplots(nrows=3, ncols=3, figsize=(8, 8))
fig.suptitle('Model Predictions')
for i, (image_path, label) in enumerate(test_images):
image = Image.open(image_path)
image = image.resize((224, 224))
image = data_trainsforms['test'](image)
image = image.unsqueeze(0)
if Use_gpu:
image = image.cuda()
output = model(image)
_, predicted = torch.max(output.data, 1)
predicted_label = example_classees[predicted.item()]
image = image.cpu().numpy()
axs[i//3, i%3].imshow(np.transpose(image.squeeze(), (1, 2, 0)))
axs[i//3, i%3].axis('off')
axs[i//3, i%3].set_title(predicted_label)
plt.show()
|
