pytorch使用在线预训练模型以及torchvison.datasets数据集进行识别分类

import torch.nn as nn
import torch
import os
import torchvision
from PIL import Image
import torch.optim as optim
from torch.utils.data import DataLoader
from torchsummary import summary
from torchvision import transforms
from tqdm import tqdm
from torchvision.models import mobilenet,MobileNet_V2_Weights
def train_data():
	# model
	device = "cuda:0"
	# model = resnet18(weights = ResNet18_Weights.IMAGENET1K_V1) # 11689512
	model = mobilenet.mobilenet_v2(weights = MobileNet_V2_Weights.IMAGENET1K_V1)
	model.fc = nn.Sequential(
		nn.Linear(1280,512),
		nn.Linear(512,66)  # 最后的分类个数根据自己数据来
	)
	batch_size = 16
	model.to(device)
	#data

	tf = transforms.Compose(
		[
			transforms.ToTensor(),
			transforms.Lambda(lambda x :x.repeat(3,1,1)),  # 针对单通道数据的扩增
			transforms.Resize((224,224)), #适应预训练模型的输入大小
			transforms.Normalize([.5,.5,.5],[.5,.5,.5])
		]
	)
	train_set = torchvision.datasets.EMNIST(root = "./data",split = "byclass",transform= tf,train = True,download = True)
	train_loader = DataLoader(train_set,batch_size = batch_size,shuffle = True)
	print(len(train_set))
	test_set = torchvision.datasets.EMNIST(root = "./data",split = "byclass",transform=tf,train = False,download = True)
	test_loader = DataLoader(test_set,batch_size = batch_size,shuffle = True)
	print(len(test_set))
	n_epoch = 100
	lr = 0.0003
	op = optim.Adam(params=model.parameters(), lr=lr, weight_decay=0.001)
	lr_scheduler = optim.lr_scheduler.StepLR(op,step_size = 1,gamma = 0.94) # 调整学习率
	criterion = torch.nn.CrossEntropyLoss()
	criterion = criterion.to(device)
	for ep in range(n_epoch):
		pbar = tqdm(train_loader)
		model.train()
		j = 0
		for x,c in pbar:
			if j >=100:
				break
			x = x.to(device)
			c = c.to(device)
			out = model(x)
			# pred = torch.argmax(out,dim = 1)
			loss = criterion(out,c)
			op.zero_grad()
			loss.backward()
			op.step()
			pbar.set_description(f"epoch:{ep},loss:{loss.item():.4f}")
			j += 1
		model.eval()
		pbar_test = tqdm(test_loader)
		correct_all = 0
		all = 0
		j = 0
		for x,c in pbar_test:
			if j >=100:
				break
			x = x.to(device)
			c = c.to(device)
			out = model(x)
			pred = torch.argmax(out,dim = 1)
			correct = (pred==c).sum().cpu().numpy()
			correct_all += correct
			all += c.shape[0]
			pbar_test.set_description(f"epoch:{ep},correct:{correct}")
			j += 1
		print(f"ep:{ep},acc:{correct_all/all : .3f},correct:{correct_all},all:{all}")
		lr_scheduler.step()
if __name__ == "__main__":
    train_data()