class CNNWrapper(PublicImageModelWrapper):
def __init__(self, labels_path):
image_shape = [32, 32, 3]
super(CNNWrapper, self).__init__(image_shape=image_shape,
labels_path=labels_path)
self.model = Net()
self.model.load_state_dict(torch.load('./checkpoint/cifar_net_best.pth')) # 获取模型参数
self.model_name = 'CNN_public'
def forward(self, x):
return self.model.forward(x)
def get_cutted_model(self, bottleneck):
return CNN_cutted(self.model, bottleneck)
WORKING_DIR = os.getcwd()
DATA_DIR = os.path.join(WORKING_DIR, "data")
TEST_DATA_DIR = os.path.join(DATA_DIR, "test")
MODELS_PATH = "S:\\models"
MODEL_FILE = "model286_tloss359.pt"
inputs, idx_to_name = load_inputs(TEST_DATA_DIR, testing=True)
# model = Net(len(inputs[0])).cuda()
model = Net(3, 20).cuda()
checkpoint = torch.load(os.path.join(MODELS_PATH, MODEL_FILE))
model.load_state_dict(checkpoint['model_state_dict'])
f_out = open(os.path.join(MODELS_PATH, MODEL_FILE.split('.')[0] + ".csv"), mode='w')
f_out.write("id,class\n")
model.eval()
with torch.no_grad():
for idx, input in enumerate(inputs):
input = input.cuda()
prediction = model.forward(input.view(1, 3, -1))
_, top_class = prediction.topk(1)
f_out.write(f"{idx_to_name[idx]},{top_class.item() + 1}\n")
f_out.close()
