TopK
@TOPK
from __future__ import print_function, division
import torch
import torch.nn as nn
import numpy as np
import torchvision
from torchvision import datasets, models, transforms
import matplotlib.pyplot as plt
import os
#import copy
#import math
import torch.nn.functional as F
plt.ion() # interactive mode
data_dir = '/home/cc/Desktop/keji/sucai/'
# Models to choose from [resnet, alexnet, vgg, squeezenet, densenet, inception]
model_name = "resnet"
# Number of classes in the dataset
num_classes = 5
# Batch size for training (change depending on how much memory you have)
batch_size = 4 #批处理尺寸(batch_size)
# Number of epochs to train for
EPOCH = 300
#pre='/home/cc/Desktop/123/train5_densenet169/net_019.pth'
# Flag for feature extracting. When False, we finetune the whole model,
# when True we only update the reshaped layer params
#feature_extract = True
feature_extract =False
# 超参数设置
pre_epoch = 0 # 定义已经遍历数据集的次数
LR = 0.001 #学习率
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
def set_parameter_requires_grad(model, feature_extracting):
if feature_extracting:
for param in model.parameters():
param.requires_grad = False
def initialize_model(model_name, num_classes, feature_extract, use_pretrained=True):
# Initialize these variables which will be set in this if statement. Each of these
# variables is model specific.
net = None
input_size = 0
if model_name == "resnet":
""" Resnet18
"""
net = models.resnet18(pretrained=use_pretrained)
set_parameter_requires_grad(net, feature_extract)
num_ftrs = net.fc.in_features
net.fc = nn.Linear(num_ftrs, num_classes)
pre='net_060.pth'
net.load_state_dict(torch.load(pre))
input_size = 224
elif model_name == "resnet34":
""" Resnet34
"""
net = models.resnet34(pretrained=use_pretrained)
set_parameter_requires_grad(net, feature_extract)
net.fc = nn.Linear(8192, num_classes)
#pre='net_015.pth'
#net.load_state_dict(torch.load(pre))
input_size = 299
elif model_name == "alexnet":
""" Alexnet
"""
net = models.alexnet(pretrained=use_pretrained)
set_parameter_requires_grad(net, feature_extract)
num_ftrs = net.classifier[6].in_features
net.classifier[6] = nn.Linear(num_ftrs,num_classes)
input_size = 224
elif model_name == "vgg":
""" VGG11_bn
"""
net = models.vgg11_bn(pretrained=use_pretrained)
set_parameter_requires_grad(net, feature_extract)
num_ftrs = net.classifier[6].in_features
net.classifier[6] = nn.Linear(num_ftrs,num_classes)
input_size = 224
elif model_name == "squeezenet":
""" Squeezenet
"""
net = models.squeezenet1_0(pretrained=use_pretrained)
set_parameter_requires_grad(net, feature_extract)
net.classifier[1] = nn.Conv2d(512, num_classes, kernel_size=(1,1), stride=(1,1))
net.num_classes = num_classes
input_size = 224
elif model_name == "resnet101":
""" Resnet101
"""
net = models.resnet101(pretrained=use_pretrained)
set_parameter_requires_grad(net, feature_extract)
num_ftrs = net.fc.in_features
net.fc = nn.Linear(num_ftrs, num_classes)
pre='/home/cc/Desktop/dj/123/train3_resnet/train_resnet101/net_018.pth'
net.load_state_dict(torch.load(pre))
input_size = 224
elif model_name == "densenet":
""" Densenet
"""
net = models.densenet121(pretrained=use_pretrained)
set_parameter_requires_grad(net, feature_extract)
num_ftrs = net.classifier.in_features
net.classifier = nn.Linear(num_ftrs, num_classes)
input_size = 224
elif model_name == "densenet169":
""" Densenet
"""
net = models.densenet169(pretrained=use_pretrained)
set_parameter_requires_grad(net, feature_extract)
num_ftrs = net.classifier.in_features
net.classifier = nn.Linear(num_ftrs, num_classes)
net.load_state_dict(torch.load(pre))
input_size = 224
elif model_name == "inception":
""" Inception v3
Be careful, expects (299,299) sized images and has auxiliary output
"""
net = models.inception_v3(pretrained=use_pretrained)
set_parameter_requires_grad(net, feature_extract)
# Handle the auxilary net
num_ftrs = net.AuxLogits.fc.in_features
net.AuxLogits.fc = nn.Linear(num_ftrs, num_classes)
# Handle the primary net
num_ftrs = net.fc.in_features
net.fc = nn.Linear(num_ftrs,num_classes)
input_size = 299
else:
print("Invalid model name, exiting...")
exit()
return net, input_size
# Initialize the model for this run
model_ft, input_size = initialize_model(model_name, num_classes, feature_extract, use_pretrained=True)
data_transforms = {
'test': transforms.Compose([
transforms.RandomResizedCrop(input_size),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
]),
'val': transforms.Compose([
transforms.Resize(input_size),
transforms.CenterCrop(input_size),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
]),
}
image_datasets = {x: datasets.ImageFolder(os.path.join(data_dir, x), data_transforms['test']) for x in ['test']}
#print(len(image_datasets['train']))
# Create training and validation dataloaders
dataloaders_dict = {x: torch.utils.data.DataLoader(image_datasets[x], batch_size=1, shuffle=False, num_workers=2) for x in ['test']}
'''
a=image_datasets['train'].class_to_idx
'''
a={ 0:'140112199012187729', 1:'140112199012197730', 2:'140112199012197731', 3:'140112199012197732', 4:'140112199012197733'}
o=0
oi=0
#checkpoint = torch.load(pre)
#model.class_to_idx = checkpoint['class_to_idx']
#model.idx_to_class = checkpoint['idx_to_class']
topk=3
for phase in ['test']:
model_ft.eval()
model_ft = model_ft.to(device)
for data in dataloaders_dict[phase]:
images, labels = data
images, labels = images.to(device), labels.to(device)
outputs = model_ft(images)
m=nn.Softmax()
ps = m(outputs)
#ps = torch.exp(outputs)
# Find the topk predictions
topk, topclass = ps.topk(3, dim=1)
#print(topk,topclass)
top_classes = [a[class_] for class_ in topclass.cpu().numpy()[0]]
#print(top_classes)
top_p = topk.cpu().detach().numpy()
print(top_p,top_classes)