Python Pytorch学习之图像检索实践_python

class TripletData(Dataset):      def __init__(self, path, transforms, split="train"):             self.path = path          self.split = split    # train or valid          self.cats = 102       # number of categories          self.transforms = transforms                   def __getitem__(self, idx):             # our positive class for the triplet          idx = str(idx%self.cats + 1)             # choosing our pair of positive images (im1, im2)          positives = os.listdir(os.path.join(self.path, idx))          im1, im2 = random.sample(positives, 2)             # choosing a negative class and negative image (im3)          negative_cats = [str(x+1) for x in range(self.cats)]          negative_cats.remove(idx)          negative_cat = str(random.choice(negative_cats))          negatives = os.listdir(os.path.join(self.path, negative_cat))             im3 = random.choice(negatives)             im1,im2,im3 = os.path.join(self.path, idx, im1), os.path.join(self.path, idx, im2), os.path.join(self.path, negative_cat, im3)             im1 = self.transforms(Image.open(im1))             im2 = self.transforms(Image.open(im2))             im3 = self.transforms(Image.open(im3))             return [im1, im2, im3]               # we'll put some value that we want since there can be far too many triplets possible      # multiples of the number of images/ number of categories is a good choice      def __len__(self):          return self.cats*8  # Transforms  train_transforms = transforms.Compose([      transforms.Resize((224,224)),      transforms.RandomHorizontalFlip(),      transforms.ToTensor(),      transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),  ])  val_transforms = transforms.Compose([      transforms.Resize((224, 224)),      transforms.ToTensor(),      transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),  ])  # Datasets and Dataloaders  train_data = TripletData(PATH_TRAIN, train_transforms)  val_data = TripletData(PATH_VALID, val_transforms)  train_loader = torch.utils.data.DataLoader(dataset = train_data, batch_size=32, shuffle=True, num_workers=4)  val_loader = torch.utils.data.DataLoader(dataset = val_data, batch_size=32, shuffle=False, num_workers=4)

class TripletLoss(nn.Module):      def __init__(self, margin=1.0):                    super(TripletLoss, self).__init__()          self.margin = margin                          def calc_euclidean(self, x1, x2):          return (x1 - x2).pow(2).sum(1)                  # Distances in embedding space is calculated in euclidean      def forward(self, anchor, positive, negative):                    distance_positive = self.calc_euclidean(anchor, positive)                    distance_negative = self.calc_euclidean(anchor, negative)                    losses = torch.relu(distance_positive - distance_negative + self.margin)                    return losses.mean()             device = 'cuda'     # Our base model  model = models.resnet18().cuda()  optimizer = optim.Adam(model.parameters(), lr=0.001)  triplet_loss = TripletLoss()     # Training  for epoch in range(epochs):            model.train()      epoch_loss = 0.0            for data in tqdm(train_loader):                    optimizer.zero_grad()          x1,x2,x3 = data          e1 = model(x1.to(device))          e2 = model(x2.to(device))          e3 = model(x3.to(device))                     loss = triplet_loss(e1,e2,e3)          epoch_loss += loss          loss.backward()          optimizer.step()                print("Train Loss: {}".format(epoch_loss.item()))                 class TripletLoss(nn.Module):      def __init__(self, margin=1.0):                    super(TripletLoss, self).__init__()          self.margin = margin                          def calc_euclidean(self, x1, x2):          return (x1 - x2).pow(2).sum(1)                  # Distances in embedding space is calculated in euclidean      def forward(self, anchor, positive, negative):                    distance_positive = self.calc_euclidean(anchor, positive)                    distance_negative = self.calc_euclidean(anchor, negative)                    losses = torch.relu(distance_positive - distance_negative + self.margin)                    return losses.mean()             device = 'cuda'        # Our base model  model = models.resnet18().cuda()  optimizer = optim.Adam(model.parameters(), lr=0.001)  triplet_loss = TripletLoss()        # Training  for epoch in range(epochs):      model.train()      epoch_loss = 0.0      for data in tqdm(train_loader):             optimizer.zero_grad()                    x1,x2,x3 = data                    e1 = model(x1.to(device))          e2 = model(x2.to(device))          e3 = model(x3.to(device))                     loss = triplet_loss(e1,e2,e3)          epoch_loss += loss          loss.backward()          optimizer.step()                print("Train Loss: {}".format(epoch_loss.item()))

#!pip install faiss-gpu  import faiss                              faiss_index = faiss.IndexFlatL2(1000)   # build the index     # storing the image representations  im_indices = []     with torch.no_grad():      for f in glob.glob(os.path.join(PATH_TRAIN, '*/*')):                    im = Image.open(f)          im = im.resize((224,224))          im = torch.tensor([val_transforms(im).numpy()]).cuda()                preds = model(im)          preds = np.array([preds[0].cpu().numpy()])          faiss_index.add(preds) #add the representation to index          im_indices.append(f)   #store the image name to find it later on               # Retrieval with a query image  with torch.no_grad():      for f in os.listdir(PATH_TEST):                    # query/test image          im = Image.open(os.path.join(PATH_TEST,f))          im = im.resize((224,224))          im = torch.tensor([val_transforms(im).numpy()]).cuda()                test_embed = model(im).cpu().numpy()                    _, I = faiss_index.search(test_embed, 5)          print("Retrieved Image: {}".format(im_indices[I[0][0]]))