jinyucai95 / edesc-pytorch Goto Github PK

@torch.no_grad()
def CIFAR10_features():
    cifar10 = torchvision.datasets.CIFAR10(root='/home/mxh/datasets/CIFAR10',train=True,download=False)
    model = torchvision.models.resnet50(weights='DEFAULT')
    model.eval()
    model.fc= torch.nn.Identity()
    to_tensor = transforms.ToTensor()
    img_list = []
    label_list = []
    for idx in tqdm(range(len(cifar10))):
        img,label = cifar10[idx]
        img = to_tensor(img).unsqueeze(0)
        img_list.append(model(img))
        label_list.append(label)
    cifar10 = torchvision.datasets.CIFAR10(root='/home/mxh/datasets/CIFAR10',train=False,download=False)
    for idx in tqdm(range(len(cifar10))):
        img,label = cifar10[idx]
        img = to_tensor(img).unsqueeze(0)
        img_list.append(model(img))
        label_list.append(label)
    img_list = torch.cat(img_list, dim=0).numpy()
    label_list = np.array(label_list)
    data = {'data':img_list, 'label':label_list}
    np.save("/home/mxh/codes/EDESC/data/CIFAR10/cifar.npy", data)