def imgClassify(inputImg): #加载训练好的模型 model = torch.load('resnet.pkl') #print(model) #固定模型参数 model.eval() #输入格式调整 BATCH_SIZE = 1 my_data = DataOperation.MyDataset(inputImg, transform=transforms.ToTensor()) my_loader = DataOperation.Data.DataLoader(dataset=my_data, batch_size=BATCH_SIZE) for batch_index, (test_x, test_y) in enumerate(my_loader): test_output = model(test_x) pred_y = torch.max(test_output, 1)[1].data.numpy() classfy = pred_y return classfy
import torch.nn as nn import torch.optim as optim import numpy as np #myfunction import DataOperation use_cuda = torch.cuda.is_available() # Hyper Parameters EPOCH = 50 # train the training data n times, to save time, we just train 1 epoch BATCH_SIZE = 50 LR = 0.001 # learning rate # 根据自己定义的那个MyDataset来创建数据集!注意是数据集!而不是loader迭代器 train_data = DataOperation.MyDataset('./StampDB/', 'train.txt', transform=transforms.ToTensor()) test_data = DataOperation.MyDataset('./StampDB/', 'test.txt', transform=transforms.ToTensor()) # valid_data = DataOperation.MyDataset('./StampDB/', 'valid.txt', transform=transforms.ToTensor()) train_loader = Data.DataLoader(dataset=train_data, batch_size=BATCH_SIZE, shuffle=True) test_loader = Data.DataLoader(dataset=test_data, batch_size=BATCH_SIZE) # valid_loader = Data.DataLoader(dataset=valid_data,batch_size=BATCH_SIZE) model = models.resnet50(pretrained=False) #if torch 0.4.2 adp = torch.nn.AdaptiveAvgPool2d(list(np.array([1, 1])))