def eval(args):
device = torch.device(f"cuda:{args.device_id}")
model = AlexNet(n_cls = 100)
model.to(device)
model.load_state_dict(torch.load(args.pretrained_path))
model.eval()
test_loader = getLoaders(split="eval", batch_size = args.batch_size, num_workers=args.num_workers )
pred_arr = []
label_arr = []
with torch.no_grad():
for idx, (img, label) in tqdm(enumerate(test_loader),total= len(test_loader)):
img = img.to(device)
pred = model.pred(img)
# mean of softmax prob from 10 different aug
pred = pred.view(-1, 10, 100)
pred = pred.mean(dim = 1)
pred_arr.append(pred.detach().cpu().numpy())
label_arr.append(label.detach().numpy())
pred_np = np.concatenate(pred_arr)
label_np = np.concatenate(label_arr)
top_1 = utils.top_k_acc(k = 1, pred = pred_np, label= label_np)
top_5 = utils.top_k_acc(k = 5, pred = pred_np, label= label_np)
confusion = utils.confusion_matrix(100, pred_np, label_np)
torch.save({
"top_1": top_1,
"top_5": top_5,
"confusion": confusion,
}, "result.pth")
print(f"top_1: {top_1*100:.2f}, top_5: {top_5*100:.2f}")
def test(imageFolder): #测试部分
is_paramatter = False #置False为导入整个模型，置True为导入参数文件
if (is_paramatter):
net = AlexNet()
model = torch.load('./model_parameter.pth',
map_location=torch.device(device)) #模型参数文件
net.load_state_dict(model)
else:
net = torch.load('./model.pkl', map_location=torch.device(device))
net = net.to(device)
torch.set_grad_enabled(False)
torch.no_grad()
net.eval()
data_num = MyDataSet(imageFolder).__len__()
for i in range(data_num):
img, ori, name = MyDataSet(imageFolder, data_transform).__getitem__(i)
out = net(img.to(device, torch.float))
predict = out.argmax(dim=1) #预测的label
probability = out[:, predict] #该label的概率
s = 'Predict result: This is a '
if predict == 0:
s += 'CAT'
else:
s += 'DOG'
s += ' with the probability of '
s += str(round(float(probability), 4))
plt.title(s)
plt.imshow(ori)
plt.savefig("./result/" + name.replace('.jpg', '') + ".png",
dpi=300) #将结果保存在result文件夹内
plt.show() #显示图片
print(name + ' Success!')
def test():
if torch.cuda.device_count() > 1:
model = torch.nn.parallel.DataParallel(
AlexNet(num_classes=opt.num_classes))
else:
model = AlexNet(num_classes=opt.num_classes)
model.load_state_dict(
torch.load(MODEL_PATH, map_location=lambda storage, loc: storage))
model.to(device)
# init value
correct1 = 0.
correct5 = 0.
total = len(test_dataloader.dataset)
with torch.no_grad():
for i, data in enumerate(test_dataloader):
# get the inputs; data is a list of [inputs, labels]
inputs, targets = data
inputs = inputs.to(device)
targets = targets.to(device)
outputs = model(inputs)
# cal top 1 accuracy
prec1 = outputs.argmax(dim=1)
correct1 += torch.eq(prec1, targets).sum().item()
# cal top 5 accuracy
maxk = max((1, 2))
targets_resize = targets.view(-1, 1)
_, prec5 = outputs.topk(maxk, 1, True, True)
correct5 += torch.eq(prec5, targets_resize).sum().item()
return correct1 / total, correct5 / total
def main():
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
data_transform = transforms.Compose( # 首先对图片预处理： resize+toTensor+normalize
[
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
# load image
img_path = ".\\pytorch_classification\\Test2_alexnet.\\dandelion.jpg" # 载入图片
assert os.path.exists(img_path), "file: '{}' dose not exist.".format(
img_path)
img = Image.open(img_path)
plt.imshow(img)
# [N, C, H, W]
img = data_transform(img)
# expand batch dimension
img = torch.unsqueeze(img, dim=0) # ^ .unsqueeze 扩充一个维度（batch维度）
# ^ .squeeze() 对数据的维度进行压缩，去掉维数为1的的维度
# read class_indict
json_path = './class_indices.json' # 读取json文件，即索引对应的类别名称
assert os.path.exists(json_path), "file: '{}' dose not exist.".format(
json_path)
json_file = open(json_path, "r")
class_indict = json.load(json_file)
# create model
model = AlexNet(num_classes=5).to(device) # 初始化网络并放入设备中
# load model weights
weights_path = "./AlexNet.pth"
assert os.path.exists(weights_path), "file: '{}' dose not exist.".format(
weights_path)
model.load_state_dict(torch.load(weights_path)) # 载入权重
model.eval() # 进入测试模式（包含dropout操作）
with torch.no_grad(): # with torch.no_grad() 禁止参数跟踪：验证中不计算损失梯度
# predict class
output = torch.squeeze(model(
img.to(device))).cpu() # torch.squeeze压缩了维数为1的的维度（batch维度）
predict = torch.softmax(output, dim=0) # 通过softmax变成概率分布
predict_cla = torch.argmax(predict).numpy() # 获得概率最大的那个索引值，并将其转化为numpy
print_res = "class: {} prob: {:.3}".format(
class_indict[str(predict_cla)], # 打印类别名称和预测概率
predict[predict_cla].numpy())
plt.title(print_res)
print(print_res)
plt.show()
def main():
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
data_transform = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
# load image
img_path = "../tulip.jpg"
assert os.path.exists(img_path), "file: '{}' dose not exist.".format(
img_path)
img = Image.open(img_path)
plt.imshow(img)
# [N, C, H, W]
img = data_transform(img)
# expand batch dimension
img = torch.unsqueeze(img, dim=0)
# read class_indict
json_path = './class_indices.json'
assert os.path.exists(json_path), "file: '{}' dose not exist.".format(
json_path)
json_file = open(json_path, "r")
class_indict = json.load(json_file)
# create model
model = AlexNet(num_classes=5).to(device)
# load model weights
weights_path = "./AlexNet.pth"
assert os.path.exists(weights_path), "file: '{}' dose not exist.".format(
weights_path)
model.load_state_dict(torch.load(weights_path))
model.eval()
with torch.no_grad():
# predict class
output = torch.squeeze(model(img.to(device))).cpu()
predict = torch.softmax(output, dim=0)
predict_cla = torch.argmax(predict).numpy()
print_res = "class: {} prob: {:.3}".format(
class_indict[str(predict_cla)], predict[predict_cla].numpy())
plt.title(print_res)
for i in range(len(predict)):
print("class: {:10} prob: {:.3}".format(class_indict[str(i)],
predict[i].numpy()))
plt.show()
def main(args=None):
logging.basicConfig(format='%(levelname)s:%(message)s',
level=logging.DEBUG)
logging.info('Arquitetura AlexNet')
parser = argparse.ArgumentParser()
parser.add_argument('--num_classes',
help='Num. de classes.',
type=int,
default=1000)
parser.add_argument('--pretrained',
help='Serão utilizados pesos pré-treinados.',
type=bool,
default=True)
parser.add_argument('--model_url',
help='Caminho para os pesos.',
default="./pesos/alexnet-owt-4df8aa71.pth")
opt = parser.parse_args(args)
# Dados
proc = Preprocessador()
imagem_url = "./imagens/raposa.jpg"
imagem = Image.open(imagem_url)
imagem = proc.executa(imagem)
#https://jhui.github.io/2018/02/09/PyTorch-Basic-operations/
imagem = imagem.unsqueeze(0)
# Instancia do modelo
model = AlexNet(opt.num_classes)
model.eval()
# Caso deseje utilizar os pesos pré-treinados
if opt.pretrained:
checkpoint = torch.load(opt.model_url)
model.load_state_dict(checkpoint)
# Utiliza a GPU se existir no computador
if torch.cuda.is_available():
model.to('cuda')
with torch.no_grad():
saida = model(imagem)
# Obtem o indice melhor ranqueado
index = np.argmax(saida[0]).item()
acuracia = torch.max(saida).item()
print(getLabel(index), acuracia)
def main():
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
data_transform = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
# /home/xulei/数据集大本营/5_flower_data
img_path = "/home/xulei/数据集大本营/5_flower_data/predict/rose01.jpeg"
assert os.path.exists(img_path), "file: '{}' dose not exist.".format(
img_path)
img = Image.open(img_path)
plt.imshow(img)
img = data_transform(img)
img = torch.unsqueeze(img, dim=0)
json_path = './class_idices.json'
assert os.path.exists(json_path), "file '{}' dose not exist.".format(
json_path)
json_file = open(json_path, 'r')
class_indict = json.load(json_file)
model = AlexNet(num_classes=5).to(device)
# load model weights
weights_path = "./AlexNet.pth"
assert os.path.exists(weights_path), "file '{}' dose not exists.".format(
weights_path)
model.load_state_dict(torch.load(weights_path))
model.eval()
with torch.no_grad():
output = torch.squeeze(model(img.to(device))).cpu()
predict = torch.softmax(output, dim=0)
predict_cla = torch.argmax(predict).numpy()
print_res = "class: {} prob: {:.3f}".format(class_indict[str(predict_cla)],
predict[predict_cla].numpy())
plt.title(print_res)
print(print_res)
plt.show()
def loadModel(self):
"""载入指定的模型"""
default_dir = os.getcwd()
modelPath = askopenfilename(
title='选择一个模型文件',
initialdir=(os.path.expanduser(default_dir)),
filetypes=[('pth文件', '*.pth'), ('All Files', '*')])
if modelPath == "":
return
try:
self.label_info.config(text="载入模型中……")
model = AlexNet(num_classes=5)
model.load_state_dict(torch.load(modelPath))
model.eval()
self.model = model
except Exception as e:
self.label_info.config(text="模型载入出错")
finally:
self.button_loadImage.config(state=tk.NORMAL)
self.label_info.config(text="请打开一张图片")
def get_prediction(image_bytes):
# 异常处理：防止传入非图片的东西
try:
weights_path = "./Alexnet.pth"
class_json_path = "./class_indices.json"
assert os.path.exists(weights_path), "weights path does not exist..."
assert os.path.exists(class_json_path), "class json path does not exist..."
# select device
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(device)
# create model
model = AlexNet(num_classes=5)
# load model weights
model.load_state_dict(torch.load(weights_path, map_location=device))
model.to(device)
model.eval()
# load class info
json_file = open(class_json_path, 'rb')
class_indict = json.load(json_file)
tensor = transform_image(image_bytes=image_bytes)
outputs = torch.softmax(model.forward(tensor).squeeze(), dim=0)
# detach去除梯度信息
prediction = outputs.detach().cpu().numpy()
# < 左对齐
template = "class:{:<15} probability:{:.3f}"
index_pre = [(class_indict[str(index)], float(p)) for index, p in enumerate(prediction)]
# sort probability
index_pre.sort(key=lambda x: x[1], reverse=True)
text = [template.format(k, v) for k, v in index_pre]
return_info = {"result": text}
except Exception as e:
return_info = {"result": [str(e)]}
return return_info
outf = f'logs'
while os.path.exists(outf):
outf += '_'
os.mkdir(outf)
batch_size = 128
learning_rate = 0.01
epochs = 100
start_epoch = 1
check_point = ''
train_loader, test_loader = load_dataset(batch_size)
model = AlexNet().to(device)
if check_point:
model.load_state_dict(torch.load(check_point))
optimizer = optim.Adam(model.parameters(), lr=learning_rate, betas=(0.9, 0.99))
criterion = nn.CrossEntropyLoss()
writer = SummaryWriter(outf + '/exp')
best_acc = float('-inf')
best_since_last = 0
for ep in range(start_epoch, epochs + 1):
if best_since_last == 20: break
elif best_since_last % 8 == 0 and best_since_last != 0:
adjust_learning_rate(optimizer, 0.5)
train_metrics = train(model, train_loader, optimizer, criterion, ep)
test_metrics = test(model, test_loader, criterion)
# print(train_metrics)
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
# 载入图片
img = Image.open("./pgy.jpg")
plt.imshow(img)
img = data_transform(img) # [H, W, Channel] --> [Channel, H, W]
img = torch.unsqueeze(img, dim=0) # [Channel, H, W] --> [1, Channel, H, W]
# 读取分类标签文本
try:
json_file = open('./class_indices.json', 'r')
class_indict = json.load(json_file)
except Exception as e:
print(e)
exit(-1)
model = AlexNet(num_classes=5)
model_weight_path = "./AlexNet.pth"
model.load_state_dict(torch.load(model_weight_path))
model.eval() # 此模式dropout使模型中的操作失效
with torch.no_grad(): # 只预测是不需要反向传播的，也就不需要梯度数据，使用no_grad()不生成梯度数据可以节约内存
output = torch.squeeze(model(img)) # 去掉batch维度 （这里batch=1）
predict = torch.softmax(output, dim=0) # 变成概率分布
predict_cla = torch.argmax(predict).numpy() # 取最大值那个位置的索引
print(class_indict[str(predict_cla)], predict[predict_cla].item())
plt.show()
def train(pertrained=False, resume_file=None):
if pertrained:
from model import alexnet
net = alexnet(pretrained=True, num_classes=NUMBER_CLASSES)
else:
from model import AlexNet
net = AlexNet(num_classes=NUMBER_CLASSES)
valid_precision = 0
policies = net.parameters()
optimizer = optim.SGD(policies,
lr=LR,
momentum=MOMENTUM,
weight_decay=WEIGHT_DECAY)
train_log = open(
"logs/train_logs_{}.log".format(
time.strftime("%Y-%m-%d-%H:%M:%S", time.localtime())), "w")
valid_log = open(
"logs/valid_logs_{}.log".format(
time.strftime("%Y-%m-%d-%H:%M:%S", time.localtime())), "w")
train_log.write("{}\t{}\t{}\n".format("epoch", "losses ", "correct"))
valid_log.write("{}\t{}\t{}\n".format("epoch", "losses ", "correct"))
# 恢复训练
if resume_file:
if os.path.isfile(resume_file):
print(("=> loading checkpoint '{}'".format(resume_file)))
checkpoint = torch.load(resume_file)
start_epoch = checkpoint['epoch']
net.load_state_dict(checkpoint['model_state_dict'])
print(("=> loaded checkpoint '{}' (epoch {})".format(
resume_file, checkpoint['epoch'])))
else:
start_epoch = 0
print(("=> no checkpoint found at '{}'".format(resume_file)))
# valid_precision = valid(net)
for epoch in range(start_epoch, EPOCHES):
batch_time = AverageMeter()
data_time = AverageMeter()
losses = AverageMeter()
correct = AverageMeter()
end = time.time()
optimizer = adjust_learning_rate(optimizer, epoch, LR, LR_steps,
WEIGHT_DECAY)
for i_batch, sample_batched in enumerate(train_dataloader):
# measure data loading time
data_time.update(time.time() - end)
inputs, labels = sample_batched
if CUDA_AVALIABLE:
outputs = net.forward(inputs.cuda())
labels = labels.long().flatten().cuda()
else:
outputs = net.forward(inputs)
labels = labels.long().flatten()
outputs = outputs.reshape([-1, NUMBER_CLASSES])
loss = criterion(outputs, labels)
# 更新统计数据
losses.update(loss.item(), inputs.size(0))
_, predicted = torch.max(outputs.data, 1)
# 计算准确率
correct.update(
(predicted == labels.long()).sum().item() / len(labels),
inputs.size(0))
optimizer.zero_grad()
loss.backward()
optimizer.step()
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
if i_batch % 10 == 0:
print(('Epoch: [{0}][{1}/{2}], lr: {lr:.5f}\t'
'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
'Data {data_time.val:.3f} ({data_time.avg:.3f})\t'
'Loss {loss.val:.4f} ({loss.avg:.4f})\t'
'Prec@1 {top1.val:.3f} ({top1.avg:.3f})\t'.format(
epoch,
i_batch,
len(train_dataloader),
batch_time=batch_time,
data_time=data_time,
loss=losses,
top1=correct,
lr=optimizer.param_groups[-1]['lr'])))
train_log.write("{:5d}\t{:.5f}\t{:.5f}\n".format(
epoch, losses.avg, correct.avg))
train_log.flush()
if epoch % 1 == 0:
valid_precision = valid(net, epoch, valid_log)
# 保存网络
if (epoch > 0 and epoch % 10 == 0) or epoch == EPOCHES - 1:
save_path = os.path.join(
"models",
"{:d}_{}_{:d}_{:d}_{:.5f}.pt".format(int(time.time()),
"alexnet", epoch,
BATCHSIZE,
valid_precision))
print("[INFO] Save weights to " + save_path)
torch.save(
{
'epoch': epoch,
'model_state_dict': net.state_dict(),
'optimizer_state_dir': optimizer.state_dict,
'loss': loss
}, save_path)
train_log.close()
valid_log.close()
def main():
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
data_transform = transforms.Compose([
transforms.Resize((360, 360)),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
# load image
#img_path = "./input/test_OUT.png"
#assert os.path.exists(img_path), "file: '{}' dose not exist.".format(img_path)
#bnyImage = Image.open(img_path)
#bnyImage = 255*np.array(bnyImage).astype('uint8')
#img = cv2.cvtColor(np.array(bnyImage), cv2.COLOR_GRAY2BGR)
#img = Image.fromarray(img)
# load image
img_path = "./input/test_OUT.png"
assert os.path.exists(img_path), "file: '{}' dose not exist.".format(
img_path)
bnyImage = Image.open(img_path)
img = cv2.cvtColor(np.array(bnyImage), cv2.COLOR_GRAY2BGR)
# one channel
img_gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
ret, img_bin = cv2.threshold(img_gray, 100, 255, cv2.THRESH_BINARY)
contours, hierarchy = cv2.findContours(img_bin, cv2.RETR_LIST,
cv2.CHAIN_APPROX_SIMPLE)
# find the left point
a1 = contours[0]
a2 = contours[1]
left1 = tuple(a1[:, 0][a1[:, :, 0].argmin()])
left2 = tuple(a2[:, 0][a2[:, :, 0].argmin()])
if left1 < left2:
leftMin = left1[0] - 25
else:
leftMin = left2[0] - 25
if leftMin > 240:
leftMin = 240
img = img[0:360, leftMin:leftMin + 360]
#print(img.shape)
img = 255 * np.array(img).astype('uint8')
img = Image.fromarray(img)
plt.imshow(img)
# [N, C, H, W]
img = data_transform(img)
# expand batch dimension
img = torch.unsqueeze(img, dim=0)
# read class_indict
json_path = './classification/male_class.json'
assert os.path.exists(json_path), "file: '{}' dose not exist.".format(
json_path)
json_file = open(json_path, "r")
class_indict = json.load(json_file)
# create model
model = AlexNet(num_classes=3).to(device)
# load model weights
weights_path = "./models/maleAlexNet_2000.pth"
assert os.path.exists(weights_path), "file: '{}' dose not exist.".format(
weights_path)
model.load_state_dict(torch.load(weights_path))
model.eval()
with torch.no_grad():
# predict class
output = torch.squeeze(model(img.to(device))).cpu()
predict = torch.softmax(output, dim=0)
predict_cla = torch.argmax(predict).numpy()
class_indict[0] = 'mid-age'
class_indict[1] = 'old-age'
class_indict[2] = 'young'
if class_indict[str(predict_cla)] == 'yM':
print_res = "class: {}".format('young', predict[predict_cla].numpy())
elif class_indict[str(predict_cla)] == 'oM':
print_res = "class: {}".format('old-age')
else:
print_res = "class: {}".format('mid-age')
#print_res = "class: {} prob: {:.3}".format(class_indict[str(predict_cla)],
# predict[predict_cla].numpy())
plt.title(print_res)
print(print_res)
plt.show()
parser.add_argument("--features",
'-f',
help="show features type",
default=None)
args = parser.parse_args()
image_path = args.image_path
pt_file = args.model_file
# image_path = "data/test/456.jpg"
# pt_file = "models/1580889991_alexnet_10_16.pt"
labels = get_labels()
net = AlexNet(num_classes=2)
print("=> Loading model ...")
checkpoint = torch.load(pt_file)
net.load_state_dict(checkpoint['model_state_dict'])
print("=> Weights loaded")
if not args.features:
predict(image_path)
elif args.features == "all":
features_map = FeatureVisualization(net, image_path)
# # show all features
features = features_map.get_all_feature()
plt.figure()
feature_columns = math.ceil((len(features) + 1) / 3)
plt.subplot(3, feature_columns, 1)
plt.imshow(features_map.origin_img)
for index, image in enumerate(features):
plt.subplot(3, feature_columns, index + 2)
plt.imshow(image)
def Test_phase(args, k):
model1 = AlexNet()
model1.cuda()
model2 = AlexNet()
model2.cuda()
model3 = AlexNet()
model3.cuda()
model4 = AlexNet()
model4.cuda()
state_dict = torch.load('checkpoints/alexnet_63.pth')
model1.load_state_dict(state_dict)
state_dict = torch.load('checkpoints/alexnet_6240.pth')
model2.load_state_dict(state_dict)
state_dict = torch.load('checkpoints/alexnet_6215.pth')
model3.load_state_dict(state_dict)
state_dict = torch.load('checkpoints/alexnet_6085.pth')
model4.load_state_dict(state_dict)
model1 = torch.nn.DataParallel(model1, device_ids=range(torch.cuda.device_count()))
model2 = torch.nn.DataParallel(model2, device_ids=range(torch.cuda.device_count()))
model3 = torch.nn.DataParallel(model3, device_ids=range(torch.cuda.device_count()))
model4 = torch.nn.DataParallel(model4, device_ids=range(torch.cuda.device_count()))
model1.eval()
model2.eval()
model3.eval()
model4.eval()
csv = csv_write(args)
dataset = Test_Dataset(args.dpath)
test_sampler = Test_Sampler(dataset._labels, n_way=args.nway, k_shot=args.kshot, query=args.query)
test_loader = DataLoader(dataset=dataset, batch_sampler=test_sampler, num_workers=8, pin_memory=True)
print('Test start!')
for i in range(TEST_SIZE):
for episode in test_loader:
data = episode.cuda()
data_shot, data_query = data[:k], data[k:]
""" TEST Method """
""" Predict the query images belong to which classes
At the training phase, you measured logits.
The logits can be distance or similarity between query images and 5 images of each classes.
From logits, you can pick a one class that have most low distance or high similarity.
ex) # when logits is distance
pred = torch.argmin(logits, dim=1)
# when logits is prob
pred = torch.argmax(logits, dim=1)
pred is torch.tensor with size [20] and the each component value is zero to four
"""
model1.eval()
features_shot = model1(data_shot)
n_sample = int(args.query/args.nway)
features_shot_mean = torch.zeros(args.nway, features_shot.size(1)).cuda()
for j in range(int(args.nway)):
start = j*args.kshot
end = (j+1)*args.kshot
features_shot_mean[j] = features_shot[start:end].mean(dim=0)
features_query = model1(data_query)
logits1 = square_euclidean_metric(features_query, features_shot_mean)
model2.eval()
features_shot = model2(data_shot)
n_sample = int(args.query/args.nway)
features_shot_mean = torch.zeros(args.nway, features_shot.size(1)).cuda()
for j in range(int(args.nway)):
start = j*args.kshot
end = (j+1)*args.kshot
features_shot_mean[j] = features_shot[start:end].mean(dim=0)
features_query = model2(data_query)
logits2 = square_euclidean_metric(features_query, features_shot_mean)
model3.eval()
features_shot = model3(data_shot)
n_sample = int(args.query/args.nway)
features_shot_mean = torch.zeros(args.nway, features_shot.size(1)).cuda()
for j in range(int(args.nway)):
start = j*args.kshot
end = (j+1)*args.kshot
features_shot_mean[j] = features_shot[start:end].mean(dim=0)
features_query = model3(data_query)
logits3 = square_euclidean_metric(features_query, features_shot_mean)
model4.eval()
features_shot = model4(data_shot)
n_sample = int(args.query/args.nway)
features_shot_mean = torch.zeros(args.nway, features_shot.size(1)).cuda()
for j in range(int(args.nway)):
start = j*args.kshot
end = (j+1)*args.kshot
features_shot_mean[j] = features_shot[start:end].mean(dim=0)
features_query = model4(data_query)
logits4 = square_euclidean_metric(features_query, features_shot_mean)
logits = (logits1 + logits2 + logits3 + logits4)/4
lsoft = F.log_softmax(-logits, dim=1).view(args.kshot, n_sample, -1)
pred = torch.argmin(logits,dim=1)
# save your prediction as StudentID_Name.csv file
csv.add(pred)
csv.close()
print('Test finished, check the csv file!')
exit()
# create model
alexnet = AlexNet(num_classes=NUM_CLASSES)
# load pretrained model
if pretrained:
alexnet_dict = alexnet.state_dict()
# print(alexnet_dict.keys())
alexnet_pretrained = models.alexnet(pretrained=True)
pretrained_dict = alexnet_pretrained.state_dict()
pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in alexnet_dict}
# print(pretrained_dict.keys())
pretrained_dict.pop("classifier.6.weight")
pretrained_dict.pop("classifier.6.bias")
alexnet_dict.update(pretrained_dict)
alexnet.load_state_dict(alexnet_dict)
# print(alexnet_dict.keys())
print("Load from pretrained")
# Freeze parameter
if freeze_layer:
for name, value in alexnet.named_parameters():
if (name != "classifier.6.weight") and (name != "classifier.6.bias"):
value.requires_grad = False
print("Freeze layer")
# train on multiple GPUs
DEVICE_IDS = list(range(GPU_NUM))
# alexnet = alexnet.to(device)
if GPU_NUM > 1:
alexnet = torch.nn.parallel.DataParallel(alexnet, device_ids=DEVICE_IDS)
port = 'com3'
bps = 115200
flag = False # 判断串口打开是否成功
data = ''
# 自定义灰度界限，大于这个值为白色，小于这个值为黑色(与原图黑白颠倒)
threshold = 100
# 二值化处理表
table = []
for i in range(256):
if i < threshold:
table.append(1)
else:
table.append(0)
model = AlexNet()
model.load_state_dict(torch.load('./logs/mnist/best.pth'))
model = model.to(device)
try:
# 打开串口，并得到串口对象
ser = serial.Serial(port, bps)
# 判断是否打开成功
if ser.is_open:
flag = True
print('Server start!')
except Exception as e:
print("ERROR: ", e)
if flag:
while True:
time.sleep(0.5)
def val(args):
# the number of N way, K shot images
k = args.nway * args.kshot
""" TODO 1.a """
" Make your own model for Few-shot Classification in 'model.py' file."
""" TODO 1.a END """
model1 = AlexNet()
model1.cuda()
model2 = AlexNet()
model2.cuda()
model3 = AlexNet()
model3.cuda()
model4 = AlexNet()
model4.cuda()
state_dict = torch.load('checkpoints/alexnet_63.pth')
model1.load_state_dict(state_dict)
state_dict = torch.load('checkpoints/alexnet_6240.pth')
model2.load_state_dict(state_dict)
state_dict = torch.load('checkpoints/alexnet_6215.pth')
model3.load_state_dict(state_dict)
state_dict = torch.load('checkpoints/alexnet_6085.pth')
model4.load_state_dict(state_dict)
#state_dict = torch.load('')
#model.load_state_dict(state_dict)
#model1 = model.clone()
model1 = torch.nn.DataParallel(model1, device_ids=range(torch.cuda.device_count()))
model2 = torch.nn.DataParallel(model2, device_ids=range(torch.cuda.device_count()))
model3 = torch.nn.DataParallel(model3, device_ids=range(torch.cuda.device_count()))
model4 = torch.nn.DataParallel(model4, device_ids=range(torch.cuda.device_count()))
model1.eval()
model2.eval()
model3.eval()
model4.eval()
csv = csv_write(args)
# Validation data loading
val_dataset = Dataset(args.dpath, state='val')
val_sampler = Sampler(val_dataset._labels, n_way=args.nway, k_shot=args.kshot, query=args.query)
val_data_loader = DataLoader(dataset=val_dataset, batch_sampler=val_sampler, num_workers=8, pin_memory=True)
""" TODO 1.b (optional) """
" Set an optimizer or scheduler for Few-shot classification (optional) "
# Default optimizer setting
#optimizer = torch.optim.Adam(model.parameters(), lr=1e-5)
#optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9, weight_decay=5e-4)
""" TODO 1.b (optional) END """
tl = Averager() # save average loss
ta = Averager() # save average accuracy
print('test start')
test_correct = 0
test_total = 0
test_loss = 0
# validation start
print('validation start')
model1.eval()
model2.eval()
model3.eval()
model4.eval()
with torch.no_grad():
vl = Averager() # save average loss
va = Averager() # save average accuracy
for j in range(VAL_TOTAL):
for episode in val_data_loader:
data, label = [_.cuda() for _ in episode]
data_shot, data_query = data[:k], data[k:] # load an episode
label_shot, label_query = label[:k], label[k:]
label_shot = sorted(list(set(label_shot.tolist())))
label_query = label_query.tolist()
labels = []
for j in range(len(label_query)):
label = label_shot.index(label_query[j])
labels.append(label)
labels = torch.tensor(labels).cuda()
""" TODO 2 ( Same as above TODO 2 ) """
""" Train the model
Input:
data_shot : torch.tensor, shot images, [args.nway * args.kshot, 3, h, w]
be careful when using torch.reshape or .view functions
data_query : torch.tensor, query images, [args.query, 3, h, w]
labels : torch.tensor, labels of query images, [args.query]
output:
loss : torch scalar tensor which used for updating your model
logits : A value to measure accuracy and loss
"""
data, label = [_.cuda() for _ in episode] # load an episode
# split an episode images and labels into shots and query set
# note! data_shot shape is ( nway * kshot, 3, h, w ) not ( kshot * nway, 3, h, w )
# Take care when reshape the data shot
data_shot, data_query = data[:k], data[k:]
#print('label : ',label)
label_shot, label_query = label[:k], label[k:]
label_shot = sorted(list(set(label_shot.tolist())))
# convert labels into 0-4 values
label_query = label_query.tolist()
labels = []
for j in range(len(label_query)):
label = label_shot.index(label_query[j])
labels.append(label)
labels = torch.tensor(labels).cuda()
""" TODO 2 ( Same as above TODO 2 ) """
""" Make a loss function and train your own model
Input:
data_shot : torch.tensor, shot images, [args.nway * args.kshot, 3, h, w]
be careful when using torch.reshape or .view functions
(25, 3, 400, 400)
data_query : torch.tensor, query images, [args.query, 3, h, w]
(20, 3, 400, 400)
labels : torch.tensor, labels of query images, [args.query]
(20)
output:
loss : torch scalar tensor which used for updating your model
logits : A value to measure accuracy and loss
"""
model1.eval()
features_shot = model1(data_shot)
n_sample = int(args.query/args.nway)
features_shot_mean = torch.zeros(args.nway, features_shot.size(1)).cuda()
for j in range(int(args.nway)):
start = j*args.kshot
end = (j+1)*args.kshot
features_shot_mean[j] = features_shot[start:end].mean(dim=0)
features_query = model1(data_query)
logits1 = square_euclidean_metric(features_query, features_shot_mean)
model2.eval()
features_shot = model2(data_shot)
n_sample = int(args.query/args.nway)
features_shot_mean = torch.zeros(args.nway, features_shot.size(1)).cuda()
for j in range(int(args.nway)):
start = j*args.kshot
end = (j+1)*args.kshot
features_shot_mean[j] = features_shot[start:end].mean(dim=0)
features_query = model2(data_query)
logits2 = square_euclidean_metric(features_query, features_shot_mean)
model3.eval()
features_shot = model3(data_shot)
n_sample = int(args.query/args.nway)
features_shot_mean = torch.zeros(args.nway, features_shot.size(1)).cuda()
for j in range(int(args.nway)):
start = j*args.kshot
end = (j+1)*args.kshot
features_shot_mean[j] = features_shot[start:end].mean(dim=0)
features_query = model3(data_query)
logits3 = square_euclidean_metric(features_query, features_shot_mean)
model4.eval()
features_shot = model4(data_shot)
n_sample = int(args.query/args.nway)
features_shot_mean = torch.zeros(args.nway, features_shot.size(1)).cuda()
for j in range(int(args.nway)):
start = j*args.kshot
end = (j+1)*args.kshot
features_shot_mean[j] = features_shot[start:end].mean(dim=0)
features_query = model4(data_query)
logits4 = square_euclidean_metric(features_query, features_shot_mean)
logits = (logits1+logits2+logits3+logits4)/4
lsoft = F.log_softmax(-logits, dim=1).view(args.kshot, n_sample, -1)
pred = torch.argmin(logits,dim=1)
""" TODO 2 END """
acc = count_acc(logits, labels)
va.add(acc)
csv.add(pred)
proto = None; logits = None; loss = None
print('val accuracy mean : %.4f' % va.item())
print('val loss mean : %.4f' % vl.item())
# initialize loss and accuracy mean
vl = None
va = None
vl = Averager()
va = Averager()
csv.close()
print('Test finished, check the csv file!')
exit()
#plt.imshow(img)
# [N, C, H, W]
img = data_transform(img)
# expand batch dimension
img = torch.unsqueeze(img, dim=0)
# read class_indict
try:
json_file = open('./class_indices.json', 'r')
class_indict = json.load(json_file)
except Exception as e:
print(e)
exit(-1)
# create model
model = AlexNet(num_classes=5)
# load model weights
model_weight_path = "./AlexNet.pth"
#, map_location='cpu'
model.load_state_dict(torch.load(model_weight_path, map_location='cpu'))
# 关闭 Dropout
model.eval()
with torch.no_grad():
# predict class
output = torch.squeeze(model(img)) # 将输出压缩，即压缩掉 batch 这个维度
predict = torch.softmax(output, dim=0)
predict_cla = torch.argmax(predict).numpy()
print(class_indict[str(predict_cla)], predict[predict_cla].item())
plt.show()
def main():
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
data_transform = transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
# load image
img_path = 'a.jpg'
assert os.path.exists(img_path), "file: '{}' does not exit.".format(
img_path)
img = Image.open(img_path)
plt.imshow(img)
# [N, C, H, W]
img = data_transform(img)
# expand batch dimension
img = torch.unsqueeze(img, dim=0)
print(img.shape)
# read class_indict
json_path = 'class_indices.json'
assert os.path.exists(json_path), "file: '{}' does not exist.".format(
json_path)
json_file = open(json_path, 'r')
class_indict = json.load(json_file)
# print(class_indict)
# create model
model = AlexNet(num_classes=5).to(device)
# load model weights
model_weight_path = 'weights/alexnet.pth'
model.load_state_dict(torch.load(model_weight_path, map_location=device))
model.eval()
print('=================================')
dummy_input = torch.randn(1, 3, 224, 224).to(device)
torch.onnx.export(model,
dummy_input,
'alexnet.onnx',
dynamic_axes={
'image': {
0: 'B'
},
'outputs': {
0: 'B'
}
},
input_names=['image'],
output_names=['outputs'],
opset_version=12)
print('=================================')
print('---------------------------------')
traced_script_module = torch.jit.trace(model, dummy_input)
traced_script_module.save("alexnet.pt")
print('---------------------------------')
with torch.no_grad():
# predict class
import time
t1 = time.time()
model(img.to(device)).cpu()
t2 = time.time()
print('torch 推理花费{}ms'.format(1000 * (t2 - t1)))
output = torch.squeeze(model(img.to(device))).cpu()
predict = torch.softmax(output, dim=0)
predict_cla = torch.argmax(predict).numpy()
print_res = "class: {} prob: {:.3}".format(class_indict[str(predict_cla)],
predict[predict_cla].numpy())
plt.title(print_res)
print(print_res)
plt.show()
def train():
try:
os.makedirs(opt.checkpoints_dir)
except OSError:
pass
if torch.cuda.device_count() > 1:
model = torch.nn.parallel.DataParallel(
AlexNet(num_classes=opt.num_classes))
else:
model = AlexNet(num_classes=opt.num_classes)
if os.path.exists(MODEL_PATH):
model.load_state_dict(
torch.load(MODEL_PATH, map_location=lambda storage, loc: storage))
model.to(device)
################################################
# Set loss function and Adam optimizer
################################################
criterion = torch.nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=opt.lr)
for epoch in range(opt.epochs):
# train for one epoch
print(f"\nBegin Training Epoch {epoch + 1}")
# Calculate and return the top-k accuracy of the model
# so that we can track the learning process.
losses = AverageMeter()
top1 = AverageMeter()
top5 = AverageMeter()
for i, data in enumerate(train_dataloader):
# get the inputs; data is a list of [inputs, labels]
inputs, targets = data
inputs = inputs.to(device)
targets = targets.to(device)
# compute output
output = model(inputs)
loss = criterion(output, targets)
# measure accuracy and record loss
prec1, prec5 = accuracy(output, targets, topk=(1, 2))
losses.update(loss.item(), inputs.size(0))
top1.update(prec1, inputs.size(0))
top5.update(prec5, inputs.size(0))
# compute gradients in a backward pass
optimizer.zero_grad()
loss.backward()
# Call step of optimizer to update model params
optimizer.step()
print(
f"Epoch [{epoch + 1}] [{i + 1}/{len(train_dataloader)}]\t"
f"Loss {loss.item():.4f}\t"
f"Prec@1 {top1.val:.3f} ({top1.avg:.3f})\t"
f"Prec@5 {top5.val:.3f} ({top5.avg:.3f})",
end="\r")
# save model file
torch.save(model.state_dict(), MODEL_PATH)
help="name of the output csv file")
parser.add_argument('--network',
type=str,
default='AlexNet',
help="Choose the type of Network [AlexNet or ResNet18 ]")
args = parser.parse_args()
use_cuda = torch.cuda.is_available()
state_dict = torch.load(args.model)
if args.network == "AlexNet":
from model import AlexNet
model = AlexNet(num_classes=20)
model.load_state_dict(state_dict)
model.eval()
elif args.network == "ResNet18":
from model import ResNet
from model import BasicBlock
model = ResNet(BasicBlock, [2, 2, 2, 2], num_classes=20)
model.load_state_dict(state_dict)
model.eval()
elif args.network == "ResNet152":
from model import resnet152, Bottleneck, ResNet
model = ResNet(Bottleneck, [3, 8, 36, 3], num_classes=20)
model.load_state_dict(state_dict)
model.eval()
training=False),
batch_size=1,
num_workers=8)
s_loader_len, t_loader_len = len(s_loader), len(t_loader)
model = AlexNet(cudable=cuda, n_class=n_class)
if cuda:
model.cuda()
opt, opt_D = model.get_optimizer(init_lr, lr_mult, lr_mult_D)
# resume or init
if not resume == '':
pretrain = torch.load(pretrain_path)
model.load_state_dict(pretrain['model'])
opt.load_state_dict(pretrain['opt'])
opt_D.load_state_dict(pretrain['opt_D'])
epoch = pretrain['epoch'] # need change to 0 when DA
else:
model.load_state_dict(utils.load_pth_model(), strict=False)
# model.load_state_dict(utils.load_pretrain_npy(), strict=False)
epoch = 0
output(' ======= START TRAINING ======= ')
for epoch in range(epoch, 100000):
model.train()
lamb = adaptation_factor(epoch * 1.0 / max_epoch)
current_lr, _ = lr_schedule(opt, epoch,
lr_mult), lr_schedule(opt_D, epoch, lr_mult_D)
def train(args):
# the number of N way, K shot images
k = args.nway * args.kshot
""" TODO 1.a """
" Make your own model for Few-shot Classification in 'model.py' file."
# model setting
model = AlexNet()
model.cuda()
""" TODO 1.a END """
# pretrained model load
if args.restore_ckpt is not None:
state_dict = torch.load(args.restore_ckpt)
model.load_state_dict(state_dict)
model = torch.nn.DataParallel(model,
device_ids=range(torch.cuda.device_count()))
if args.test_mode == 1:
Test_phase(model, args, k)
else:
# Train data loading
dataset = Dataset(args.dpath, state='train')
train_sampler = Train_Sampler(dataset._labels,
n_way=args.nway,
k_shot=args.kshot,
query=args.query)
data_loader = DataLoader(dataset=dataset,
batch_sampler=train_sampler,
num_workers=8,
pin_memory=True)
# Validation data loading
val_dataset = Dataset(args.dpath, state='val')
val_sampler = Sampler(val_dataset._labels,
n_way=args.nway,
k_shot=args.kshot,
query=args.query)
val_data_loader = DataLoader(dataset=val_dataset,
batch_sampler=val_sampler,
num_workers=8,
pin_memory=True)
""" TODO 1.b (optional) """
" Set an optimizer or scheduler for Few-shot classification (optional) "
# Default optimizer setting
#optimizer = torch.optim.Adam(model.parameters(), lr=1e-5)
optimizer = torch.optim.SGD(model.parameters(),
lr=0.001,
momentum=0.9,
weight_decay=5e-4)
""" TODO 1.b (optional) END """
tl = Averager() # save average loss
ta = Averager() # save average accuracy
# training start
print('train start')
train_correct = 0
train_total = 0
train_loss = 0
test_correct = 0
test_total = 0
test_loss = 0
model.train()
for i in range(args.se + 1, TOTAL):
for episode in data_loader:
optimizer.zero_grad()
data, label = [_ for _ in episode] # load an episode
# split an episode images and labels into shots and query set
# note! data_shot shape is ( nway * kshot, 3, h, w ) not ( kshot * nway, 3, h, w )
# Take care when reshape the data shot
data_shot, data_query = data[:k], data[k:]
label_shot, label_query = label[:k], label[k:]
label_shot = sorted(list(set(label_shot.tolist())))
# convert labels into 0-4 values
label_query = label_query.tolist()
labels = []
for j in range(len(label_query)):
label = label_shot.index(label_query[j])
labels.append(label)
labels = torch.tensor(labels).cuda()
""" TODO 2 ( Same as above TODO 2 ) """
""" Train the model
Input:
data_shot : torch.tensor, shot images, [args.nway * args.kshot, 3, h, w]
be careful when using torch.reshape or .view functions
data_query : torch.tensor, query images, [args.query, 3, h, w]
labels : torch.tensor, labels of query images, [args.query]
output:
loss : torch scalar tensor which used for updating your model
logits : A value to measure accuracy and loss
"""
features_shot = model(data_shot.cuda())
n_sample = int(args.query / args.nway)
features_shot_mean = torch.zeros(args.nway,
features_shot.size(1)).cuda()
for j in range(int(args.nway)):
start = j * args.kshot
end = (j + 1) * args.kshot
features_shot_mean[j] = features_shot[start:end].mean(dim=0)
features_query = model(data_query.cuda())
logits = square_euclidean_metric(features_query,
features_shot_mean)
labels_expanded = labels.view(args.query, 1, 1)
labels_expanded = labels_expanded.expand(args.query, args.nway, 1)
lsoft = F.log_softmax(-logits, dim=1).view(args.kshot, n_sample,
-1)
labels_expanded = labels_expanded.view(lsoft.size())
loss = -lsoft.gather(2, labels_expanded).squeeze().view(-1).mean()
_, pred = lsoft.max(2)
""" TODO 2 END """
acc = count_acc(logits, labels)
tl.add(loss.item())
ta.add(acc)
loss.backward()
optimizer.step()
proto = None
logits = None
loss = None
if (i + 1) % PRINT_FREQ == 0:
print('train {}, loss={:.4f} acc={:.4f}'.format(
i + 1, tl.item(), ta.item()))
# initialize loss and accuracy mean
tl = None
ta = None
tl = Averager()
ta = Averager()
# validation start
if (i + 1) % VAL_FREQ == 0:
print('validation start')
model.eval()
with torch.no_grad():
vl = Averager() # save average loss
va = Averager() # save average accuracy
for j in range(VAL_TOTAL):
for episode in val_data_loader:
data, label = [_.cuda() for _ in episode]
data_shot, data_query = data[:k], data[
k:] # load an episode
label_shot, label_query = label[:k], label[k:]
label_shot = sorted(list(set(label_shot.tolist())))
label_query = label_query.tolist()
labels = []
for j in range(len(label_query)):
label = label_shot.index(label_query[j])
labels.append(label)
labels = torch.tensor(labels).cuda()
""" TODO 2 ( Same as above TODO 2 ) """
""" Train the model
Input:
data_shot : torch.tensor, shot images, [args.nway * args.kshot, 3, h, w]
be careful when using torch.reshape or .view functions
data_query : torch.tensor, query images, [args.query, 3, h, w]
labels : torch.tensor, labels of query images, [args.query]
output:
loss : torch scalar tensor which used for updating your model
logits : A value to measure accuracy and loss
"""
optimizer.zero_grad()
data, label = [_.cuda()
for _ in episode] # load an episode
# split an episode images and labels into shots and query set
# note! data_shot shape is ( nway * kshot, 3, h, w ) not ( kshot * nway, 3, h, w )
# Take care when reshape the data shot
data_shot, data_query = data[:k], data[k:]
label_shot, label_query = label[:k], label[k:]
label_shot = sorted(list(set(label_shot.tolist())))
# convert labels into 0-4 values
label_query = label_query.tolist()
labels = []
for j in range(len(label_query)):
label = label_shot.index(label_query[j])
labels.append(label)
labels = torch.tensor(labels).cuda()
""" TODO 2 ( Same as above TODO 2 ) """
""" Make a loss function and train your own model
Input:
data_shot : torch.tensor, shot images, [args.nway * args.kshot, 3, h, w]
be careful when using torch.reshape or .view functions
(25, 3, 400, 400)
data_query : torch.tensor, query images, [args.query, 3, h, w]
(20, 3, 400, 400)
labels : torch.tensor, labels of query images, [args.query]
(20)
output:
loss : torch scalar tensor which used for updating your model
logits : A value to measure accuracy and loss
"""
features_shot = model(data_shot.cuda())
n_sample = int(args.query / args.nway)
features_shot_mean = torch.zeros(
args.nway, features_shot.size(1)).cuda()
for j in range(int(args.nway)):
start = j * args.kshot
end = (j + 1) * args.kshot
features_shot_mean[j] = features_shot[
start:end].mean(dim=0)
features_query = model(data_query.cuda())
logits = square_euclidean_metric(
features_query, features_shot_mean)
labels_expanded = labels.view(args.query, 1, 1)
labels_expanded = labels_expanded.expand(
args.query, args.nway, 1)
lsoft = F.log_softmax(-logits, dim=1).view(
args.kshot, n_sample, -1)
labels_expanded = labels_expanded.view(lsoft.size())
loss = -lsoft.gather(
2, labels_expanded).squeeze().view(-1).mean()
_, pred = lsoft.max(2)
""" TODO 2 END """
acc = count_acc(logits, labels)
vl.add(loss.item())
va.add(acc)
proto = None
logits = None
loss = None
print('val accuracy mean : %.4f' % va.item())
print('val loss mean : %.4f' % vl.item())
# initialize loss and accuracy mean
vl = None
va = None
vl = Averager()
va = Averager()
if (i + 1) % SAVE_FREQ == 0:
PATH = 'checkpoints/%d_%s.pth' % (i + 1, args.name)
torch.save(model.module.state_dict(), PATH)
print('model saved, iteration : %d' % i)
import torchvision.transforms as transforms
from PIL import Image
from model import AlexNet
import json
transform = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
try:
json_file = open('./class_indices.json', 'r')
classes = json.load(json_file)
except Exception as e:
print(e)
exit(-1)
net = AlexNet(num_classes=5, init_weights=True)
net.load_state_dict(torch.load('./AlexNet.pth'))
net.eval()
im = Image.open('Download.jpeg') # [H,W,C]
im = transform(im) # [C,H,W]
im = torch.unsqueeze(im, dim=0) # [N,C,H,W] 给tensor添加一个维度：batchsize
with torch.no_grad():
outputs = net(im)
predict = torch.max(outputs, dim=1)[1].data.numpy() # [4]
print(classes[str(predict[0])])
