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])))
