2023_12_15

    def test_img(linear, net, imgs,files_name):
        #test doesn't need cuda
        net.eval()
        linear.eval()
        my_img_test={'file_name':[],'predict':[],'True?':[]}
        classes = ["airplane", "automobile", "bird", "cat", "deer", "dog", "frog", "horse", "ship", "truck"]
        for i,img in enumerate(imgs):
            # print(img.shape)
            img=img.cuda(non_blocking=True)
            feature = net(img) #minibatch is not used

            predict = linear(feature)
            max_predict=torch.argmax(predict)
            class_name=classes[max_predict]
            ture_or_false=is_subset(class_name,files_name[i])
            my_img_test['file_name'].append(files_name[i])
            my_img_test['predict'].append(class_name)
            my_img_test['True?'].append(ture_or_false)
            print('predict:',class_name,'--------True?:',ture_or_false)
        my_img_test=pd.DataFrame(my_img_test)
        my_img_test.to_csv('results\\test_imgs\\my_img_test.csv')
    checkpoint = torch.load("results\\model_last.pth")
    model.load_state_dict(checkpoint['state_dict'])
    model.eval()
    linear=torch.load('results\\linear.pth')#

    # print(a)
    # linear.load('results\\linear.pth')
    linear.eval()
    # the picture file
    folder_path = 'results\\test_imgs'

    # all the picture/imgs
    imgs = []
    files_name=[]


    def is_subset(sub, main):
        sub_set = set(sub)
        main_set = set(main)

        # the sub_set is a subset of the main_set?
        return sub_set.issubset(main_set)


    for filename in os.listdir(folder_path):
        if filename.endswith('.jpg') or filename.endswith('.png'):
            img_path = os.path.join(folder_path, filename)
            files_name.append(filename)
            img = Image.open(img_path)
            img = img.resize((32, 32))

            img = test_transform(img)
            img = torch.tensor(img)
            # transverse(img) to the batch dimension
            # this resize the picture 32*32*3 to 1*32*32*3
            img = img.view(1, *img.size())
            # print(img.size())
            imgs.append(img)
    # can use
    # SplitBatchNorm=nn.BatchNorm2d
    # model.encoder_q.SplitBatchNorm=SplitBatchNorm
    
    #
    #
    test_img(linear, model.encoder_q, imgs,files_name)
复制