基础分类模型fashmnist

数据集：https://github.com/zalandoresearch/fashion-mnist/tree/master/data/fashion

import os
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader

class FMDataset(Dataset):
	def __init__(self, df, transform=None):
		super.__init__()
		self.df = df
		self.transform = transform
		self.images = df.iloc[:, 1:].values.astype(np.uint8)
		self.labels = df.iloc[:, 0].values

	def __len__(self):
		return len(self.images)

	def __getitem__(self, idx):
		image = self.images[idx].reshape(28, 28, 1) # img_sz
		label = int(self.labels[idx])
		if self.transform is not None:
			image = self.transform(image)
		else:
			image = torch.tensor(image / 255, dtype = torch.float)
		label = torch.tensor(label, dtype=torch.long)

class Net(nn.Module):
	def __init__(self):
		super(Net, self).__init__()
		self.conv = nn.Sequential(
			# 灰度图in_channel=1(RGB=3)
			nn.Conv2d(1, 32, 5),	# 卷积
			nn.ReLU(),	# 激活
			nn.MaxPool2d(2, stride=2),	# 池化
			nn.Dropout(0.3),	 # 按照0.3的概率将tensor置0
			nn.Conv2d(32, 64, 5),
			nn.ReLU(),
			nn.MaxPoll2d(2, stride=2),
			nn.Dropout(0.3),
		)
		self.fc = nn.Sequential(
			nn.Linear(64*4*4, 512),	# 线性计算
			nn.ReLU(),
			nn.Linear(512, 10)	# 输出10个分类
		)
		
	def forward(self, x):
		x = self.conv(x)
		x = x.view(-1, 64*4*4)
		x = self.fc(x)
		return x

if __name__ == "__main__":
	# 超参数
	device = torch.device("cuda:0")
	batch_size = 32
	num_workers = 4
	lr = 1e-4
	epochs = 100
	img_size = 28

	# 加载数据
	data_transform = tranform.Compose([
		transform.ToPILImage(),
		transform.Resize(img_size),
		transform.ToTensor(),
	])
	train_df = pd.read_csv("./fashion-mnist_train.csv")
	test_df = pd.read_csv("./fashion-mnist_test.csv")
	train_data = FMDataset(train_df)
	test_data = FMDataset(test_df)
	train_loader = DataLoder(train_data, batch_size=batch_size,
								shuffle=True, num_workers=num_workers, drop_last=True)
	test_loader = DataLoder(test_data, batch_size=batch_size,
								shuffle=False, num_workers=num_workers)

	# 加载模型
	model = Net()
	model = model.cuda(device)

	# 损失函数
	criterion = nn.CrossEntropyLoss()

	# 优化器，进行模型的权重偏置等参数调整
	optimizer = optim.Adam(model.parameters(), lr=lr)

	# 训练、验证（测试）
	def train():
		model.train()
		train_loss = 0
		for data, label in train_loader:
			data, label = data.to(device), label.to(device)
			optimizer.zero_grad() # 因为求导过程累加,所以这里要归0,每次都要重新计算单次的梯度
			output = model(data)	# 前向传播
			loss = criterion(output, label)	# 损失函数计算损失
			loss.backward()	# 方向传播，更新梯度
			optimizer.step()	# 优化器更新权重偏置等参数
			tran_loss += loss.item() * data.size(0)
		train_loss = train_loss / len(train_loader.dataset)
		return train_loss
	def val():
		model.eval()
		val_loss = 0
		gt_labels = []
		pred_labels = []
		with torch.no_grad():	# 不做梯度计算
			for data, label in test_loader:
				data, label = data.to(device), label.to(device)
				output = model(data)
				preds = torch.argmax(output, 1)
				gt_labels.append(label.cpu().data.numpy())
				pred_labels.append(preds.cpu().data.numpy())
				loss = cirterion(output, label)
				val_loss += loss.item() * data.size(0)
		val_loss = val_loss / len(test.loader.dataset)
		gt_labels, pred_labels = np.concatenate(gt_labels), np.concatenate(pred_labels)
		acc = np.sum(gt_labels == pred_labels) / len(pred_labels)
		return val_loss, acc

	best_acc = .0
	for epoch in range(1, epochs + 1):
		train_loss = train()
		val_loss = val()
		print("Epoch({}/{}): train_loss: {:.4f} \tval_loss: {:.4f} \tAccuracy: {:.4f}".format(
			epoch, epochs, train_loss, val_loss, acc))
		if acc > best_acc:
			torch.save(model, "fashmnist_best.pth")
			best_acc = acc
		torch.save(model, "fashmnist_last.pth")

	# 导出onnx模型
	input = torch.randn(1, 1, img_size, img_size).to(device)
	torch.onnx.export(
		model,
		input,
		"fashminist_last.onnx",
		opset_version = 11,
		do_constant_folding=True,
		input_names=["input"],
		output_names=["output"]
	)