import os
import itertools
from torchvision import transforms
from torch.utils.data import Dataset
import os
from PIL import Image
import torch
from torch.utils.data import DataLoader
import torch
import torch.nn as nn
from torchvision.datasets import mnist # 导入 pytorch 内置的 mnist 数据
from torch.autograd import Variable
from torchvision import transforms
from torch.utils.data import DataLoader
import numpy as np
import torch
from torchvision.datasets import mnist # 导入 pytorch 内置的 mnist 数据
from torch import nn
from torch.autograd import Variable
import time
class MyNet(nn.Module):
def __init__(self):
super(MyNet, self).__init__()
self.conv = nn.Sequential(
nn.Conv2d(3, 32, (3, 3), (1, 1), 2), # out (32, 482, 446)
nn.BatchNorm2d(32),
nn.ReLU(inplace = True),
nn.Conv2d(32, 64, (3, 3), (1, 1), 2), # out (64, 482, 446)
nn.BatchNorm2d(64),
nn.ReLU(inplace=True),
nn.Conv2d(64, 128, (3, 3), (1, 1), 2), # out (128, 482, 446)
nn.BatchNorm2d(128),
nn.ReLU(inplace=True),
nn.MaxPool2d(2, 2) # out (128, 241, 223)
)
self.fc = nn.Sequential(
nn.Linear(7058432, 300),
nn.ReLU(True),
nn.Linear(300, 256),
nn.ReLU(True),
nn.Linear(256, 3),
nn.ReLU(True)
)
def forward(self, x):
# in torch.Size([batchsize, 1, 28, 28])
x = self.conv(x) # out torch.Size([batchsize, 128, 14, 14])
x = x.view(x.size(0), -1) # out (batchsize, 128*14*14)
x = self.fc(x) # out (batchsize, 10)
return x
def concat(s, l):
return([s + i for i in l])
mydic = {
'embedded_small':0, 'surface_small':1, 'un_visible':2
}
class MyData():
def __init__(self, root_dir, label_dir = None, ):
path_append = []
self.img_path = []
if label_dir is None:
label_dir = ['embedded_small', 'surface_small', 'un_visible']
self.root_dir = root_dir
self.label_dir = label_dir
for i in range(len(self.label_dir)):
full_name = concat(self.root_dir + self.label_dir[i] + '/', os.listdir(os.path.join(self.root_dir + self.label_dir[i])))
self.img_path.append(full_name)
self.img_path = list(itertools.chain.from_iterable(self.img_path))
def __getitem__(self, idx):
img_name = self.img_path[idx]
img_item_path = img_name
img = Image.open(img_item_path) #PIL img
img = self.data_tf(img)
print(img.size)
img = np.asarray(img) #array
print(img.shape)
label = mydic[str.split(img_name, '/')[2]]
return img, label
def __len__(self):
return len(self.img_path)
@staticmethod
def data_tf(x):
x = transforms.CenterCrop((482, 446))(x) # PIL img
x = np.array(x, dtype='float32') / 255
x = (x - 0.5) / 0.5 # 标准化,这个技巧之后会讲到
x = x.reshape((3, 482, 446)) # 拉平
x = torch.from_numpy(x)
return x
if __name__ == '__main__':
dataset = MyData(root_dir='./dataset/')
train_data = DataLoader(dataset, batch_size=5, shuffle=True)
net = MyNet()
# 定义 loss 函数
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(net.parameters(), 1e-1) # 使用随机梯度下降,学习率 0.1
# 开始训练
losses = []
acces = []
eval_losses = []
eval_acces = []
for e in range(1):
train_loss = 0
train_acc = 0
net.train()
for im, label in train_data:
im = Variable(im)
label = Variable(label)
print(label)
# 前向传播
out = net(im)
print(out.shape)
print(label.shape)
loss = criterion(out, label)
print(label.size)
# 反向传播
optimizer.zero_grad()
loss.backward()
optimizer.step()
# 记录误差
train_loss += loss.item()
# 计算分类的准确率
pred = out.max(1)
num_correct = (pred == label).sum().item()
acc = num_correct / im.shape[0]
train_acc += acc
losses.append(train_loss / len(train_data))
acces.append(train_acc / len(train_data))
# 在测试集上检验效果
# eval_loss = 0
# eval_acc = 0
# net.eval() # 将模型改为预测模式
# for im, label in test_data:
# im = Variable(im)
# label = Variable(label)
# out = net(im)
# loss = criterion(out, label)
# print(loss.item)
# # 记录误差
# eval_loss += loss.item()
# # 记录准确率
# _, pred = out.max(1)
# num_correct = (pred == label).sum().item()
# acc = num_correct / im.shape[0]
# eval_acc += acc
# eval_losses.append(eval_loss / len(test_data))
# eval_acces.append(eval_acc / len(test_data))
print('epoch: {}, Train Loss: {:.6f}, Train Acc: {:.6f}'
.format(e, train_loss / len(train_data), train_acc / len(train_data)))
PYTORCH
最新推荐文章于 2024-10-01 12:11:59 发布