PyTorch 常用命令

example_tensor.device

example_tensor.shape
example_tensor.size(1)

example_tensor.numel()

example_scalar = example_tensor[1, 1, 0]
example_scalar.item()
example_tensor[:, 0, 0]

torch.ones_like(example_tensor)
torch.zeros_like(example_tensor)
torch.randn_like(example_tensor)
torch.randn(2, 2, device='cpu') # Alternatively, for a GPU tensor, you'd use device='cuda'
torch.arange(4 * 5 * 6).view(4, 5, 6) # shape [4,5,6]

real_clone = real.clone()

F.one_hot(labels % n_classes)

nn.init.xavier_uniform_(tensor) # fill with uniform distribution, values scaled by 'gain'
nn.init.constant_(bias, 0) # fill with constant
nn.init.orthogonal_(param) # with a (semi) orthogonal matrix

(example_tensor - 5) * 2 # 顺序有影响
torch.mul(tensor_x, tensor_x) # 逐元素积
tensor_x * tensor_y # 逐元素积 = torch.mul(x, y)

torch.matmul(tensor_x, tensor_x)  # 张量乘法，类似矩阵乘法, 支持维度的 broadcast， 当输入均为 3 维时，和 bmm 一样 
tensor_x @ tensor_y # 和 torch.matmul(A, B) 一样

torch.bmm(tensor_x, tensor_y)  # 和 torch.matmul(A, B) 一样, 但只支持三维 tensor，常用于 Attention 计算，其中, tensor_x 是 weight_matrix， tensor_y 是 hidden_matrix

torch.dot(tensor_vec_x, tensor_vec_y)  # 和 torch.matmul(A, B) 类似, 但只支持一维 tensor

example_tensor.mean()
example_tensor.std()

# mean over one dimension
example_tensor.mean(0)
example_tensor.mean(dim=0) # or axis=0
torch.mean(example_tensor, dim=0) # or axis=0

a = torch.randn(1, 3)  # [0.3618 1.2095 -0.3403]
torch.prod(a)  # -0.14892165020372308

torch.isclose(tensorA, tensorB)

torch.norm(torch.ones(4, 2), dim=1)
# tensor([2., 2.])
torch.norm(torch.ones(4, 2), dim=1)
# tensor([1.4142, 1.4142, 1.4142, 1.4142])

start = torch.arange(1., 5.)
end = torch.empty(4).fill_(10)
torch.lerp(start, end, 0.5)

tensor_A.transpose(0, 1)
tensor_A.transpose(-2, -1)

tensor_A = torch.ones(4)
tensor_A.expand([2,3,4])

a = torch.randn(2, 3)  # shape= [2, 3]
torch.unsqueeze(a, dim=0)  # shape = [1, 2, 3]
a.unsqueeze(dim=-1) # shape = [2, 3, 1]
# dim range [-a.dim() -1, a.dim()]

m = torch.zeros(2, 1, 2, 1, 2) # shape = [2, 1, 2, 1, 2]
n = torch.squeeze(m)  # shape = [2， 2， 2]

a = torch.tensor([[[1,2,3],[4,5,6]]]) # shape = [1, 2, 3]
a_p =a.permute(2,0,1)  # shape = [3, 1, 2]
# tensor([[[ 1,  4]], [[ 2,  5]], [[ 3,  6]]])

a = torch.tensor([[[1,2,3],[4,5,6]]]) # shape = [1, 2, 3]
a_v =a.view(1,3,2)  # shape = [3, 1, 2]
# tensor([[[ 1,  2], [ 3,  4], [ 5,  6]]])

# list to tensor(cpu)
l0 = [1, 2, 3]
t = torch.Tensor(l0)

# tensor(cpu) to list
l1 = t.numpy().tolist()

# Tensor(cpu) to numpy
t = torch.ones(5)
a = t.numpy()
print(a) #  [1. 1. 1. 1. 1.]
# 注意，转换后的tensor与numpy指向同一地址，所以，对一方的值改变另一方也随之改变

# Numpy to tensor(cpu)
t1 = torch.from_numpy(a)

m = nn.Identity()
m(torch.ones(2, 2))
# tensor([[1., 1.], [1., 1.]]) 

class ExampleModule(nn.Module):
    def __init__(self, input_dims, output_dims):
        super(ExampleModule, self).__init__()
        self.linear = nn.Linear(input_dims, output_dims)
        self.exponent = nn.Parameter(torch.tensor(1.))
        self._init_params()

    def forward(self, x): # What to return after pass x to model
        x = self.linear(x)
        x = x ** self.exponent 
        return x
    
    def _init_params(self):  # For special cases
        # init parameters here

example_model = ExampleModule(10, 2)
list(example_model.parameters()) # Parameters without name
list(example_model.named_parameters()) # Parameters & names

input = torch.randn(2, 10)
example_model(input)

def get_path(base_dir, base_name, suffix):
    return os.path.join(base_dir, base_name + suffix)

def save_gpt_model(model, base_dir, base_name):
    raw_model = model.module if hasattr(model, "module") else model
    torch.save(raw_model.state_dict(), get_path(base_dir, base_name, '.pt'))

def get_path(base_dir, base_name, suffix):
    return os.path.join(base_dir, base_name + suffix)

# Need config and checkpoint file
def load_model(model, model_weights_path, device, copy_to_cpu=True):
    raw_model = model.module if hasattr(model, "module") else model
    map_location = lambda storage, loc: storage if copy_to_cpu else None
    raw_model.load_state_dict(torch.load(model_weights_path, map_location))
    return raw_model.to(device)

linear = nn.Linear(10, 2)
example_input = torch.randn(3, 10)
example_output = linear(example_input)

mlp_layer = nn.Sequential(
    nn.Linear(5, 2),
    nn.BatchNorm1d(2),
    nn.ReLU()
)
test_example = torch.randn(5,5) + 1

print(mlp_layer(test_example))

embedding = nn.Embedding(10, 3)
input = torch.LongTensor([[1,2,4,5],[4,3,2,9]])
embedding(input)

batchnorm = nn.BatchNorm1d(2)
batchnorm_output = batchnorm(relu_output)

relu = nn.ReLU()
relu_output = relu(example_output)

nn.LeakyReLU(0.2)

def collate_fn(data):
    data.sort(key=lambda x: len(x), reverse=True)
    data = pad_sequence(data, batch_first=True, padding_value=0)
    return data

data_loader = DataLoader(data, batch_size=2, shuffle=True, collate_fn=collate_fn)

def collate_fn(data):
    data.sort(key=lambda x: len(x), reverse=True)
    seq_len = [s.size(0) for s in data] # 获取数据真实的长度
    data = pad_sequence(data, batch_first=True)    
    data = pack_padded_sequence(data, seq_len, batch_first=True)
    return data

data_loader = DataLoader(data, batch_size=2, shuffle=True, collate_fn=collate_fn)

def collate_fn(data):
    data.sort(key=lambda x: len(x), reverse=True)
    data = pack_sequence(data)
    return data

data_loader = DataLoader(data, batch_size=2, shuffle=True, collate_fn=collate_fn)

def collate_fn(data):
    data.sort(key=lambda x: len(x), reverse=True)
    seq_len = [s.size(0) for s in data]
    data = pad_sequence(data, batch_first=True).float()    
    data = data.unsqueeze(-1)
    data = pack_padded_sequence(data, seq_len, batch_first=True)
    return data
   
data_loader = DataLoader(data, batch_size=2, shuffle=True, collate_fn=collate_fn)

nn.Conv2d(input_channels, output_channels, kernel_size, stride)

nn.ConvTranspose2d(input_channels, output_channels, kernel_size, stride)

nn.BatchNorm2d(output_channels)

nn.Upsample((starting_size, starting_size), mode='bilinear')

nn.functional.interpolate((starting_size, starting_size), mode='bilinear')

criterion = nn.BCEWithLogitsLoss()

import logging
from tqdm import tqdm
import numpy as np
import torch
from torch.utils.data.dataloader import DataLoader
from torch.utils.tensorboard import SummaryWriter
from model import save_model

logger = logging.getLogger(__name__)

class Trainer:
    def __init__(self, model, output_dir='./', learning_rate=3e-4, betas=(0.9, 0.95), grad_norm_clip=1.0, multi_gpu=False):
        ...

    def _save_checkpoint(self, base_dir, info, valid_loss):
        ...
.        
    def train(self, train_dataset, test_dataset=None, n_epochs=10, batch_size=64, num_workers=0):
        ...

class Trainer:
    def __init__(self, model, output_dir='./', learning_rate=3e-4, betas=(0.9, 0.95), grad_norm_clip=1.0, multi_gpu=False):
        self.model = model

        self.learning_rate = learning_rate
        self.betas = betas
        self.grad_norm_clip = grad_norm_clip
        self.lr_decay = lr_decay

        self.output_dir = output_dir
        self.writer = SummaryWriter(self.output_dir)

        self.device = 'cpu'
        if torch.cuda.is_available():
            self.device = torch.cuda.current_device()
            if multi_gpu:
                self.model = torch.nn.DataParallel(self.model).to(self.device)

class Trainer:
    def _save_checkpoint(self, base_dir, info, valid_loss):
        """save checkpoint during training. Format: model_{info}_{valid_loss}"""
        base_name = f'model_{info}_{valid_loss:.3f}'
        logger.info(f'Save model {base_name}')
        save_model(self.model, base_dir, base_name)

class Trainer:
    def train(self, train_dataset, test_dataset=None, n_epochs=10, batch_size=64, num_workers=0):
        model = self.model
        raw_model = model.module if hasattr(self.model, "module") else model
        optimizer = raw_model.configure_optimizers(self.learning_rate, self.betas)

        def run_epoch(split):
            ...
            return loss

        train_loss, test_loss, best_loss = float('inf'), float('inf'), float('inf')
        for epoch in range(n_epochs):
            train_loss = run_epoch('train')

            if test_dataset is not None:
                test_loss = run_epoch('test')
                if test_loss < best_loss:
                    self._save_model(self.output_dir, str(epoch + 1), test_loss)
                    best_loss = test_loss

            if test_dataset is None:  # save every step if no test dataset
                self._save_model(self.output_dir, str(epoch + 1), train_loss)

        final_loss = test_loss if test_dataset is not None else train_loss
        self._save_model(self.output_dir, 'final', final_loss)

class Trainer:
    def train(self, train_dataset, test_dataset=None, n_epochs=10, batch_size=64, num_workers=0):
        ...
        def run_epoch(split):
            is_train = True if split == 'train' else False
            model.train(is_train)
            data = train_dataset if is_train else test_dataset
            loader = DataLoader(data, shuffle=True, pin_memory=True, batch_size=batch_size, num_workers=num_workers)

            losses = []
            pbar = tqdm(enumerate(loader), total=len(loader)) if is_train else enumerate(loader)
            for it, (x, y) in pbar:
                x = x.to(self.device)
                y = y.to(self.device)

                with torch.set_grad_enabled(is_train):
                    logits, loss = model(x, y)
                    loss = loss.mean()  # collapse all losses if they are scattered on multiple gpus
                    losses.append(loss.item())

                if is_train:
                    model.zero_grad()
                    loss.backward()
                    torch.nn.utils.clip_grad_norm_(model.parameters(), self.grad_norm_clip)
                    optimizer.step()
                    pbar.set_description(f"epoch {epoch + 1} iter {it}: train loss {loss.item():.5f}")

            loss = float(np.mean(losses))
            logger.info(f'{split}, epoch: {epoch + 1}/{n_epochs}, loss: {loss:.4f}')
            self.writer.add_scalar('loss', loss, epoch + 1)  # rewrite if having test dataset
            return loss
         ...

import torch.optim as optim
adam_opt = optim.Adam(model.parameters(), lr=1e-1)

train_example = torch.randn(100,5) + 1
adam_opt.zero_grad() # 1
cur_loss = torch.abs(1 - mlp_layer(train_example)).mean() # 2
cur_loss.backward() # 3
adam_opt.step() # 4
print(cur_loss) 

for epoch in range(N_EPOCH):
    # train model
    ...    
    # evaluate model:
    net.eval() # changing the behavior of the nn.Module
    with torch.no_grad():  # changing the behavior of the autograd to disable gradient
        y_pred = net(x.float())
        mseLoss = torch.nn.MSELoss(reduction='sum')(y_pred, y.float())
        loss_values.append(loss_v)

    # enable training
    net.train()

@torch.no_grad()
def eval(model, data):
  model.eval()
  # Rest of eval code

model = [Parameter(torch.randn(2, 2, requires_grad=True))]
optimizer = SGD(model, 0.1)
scheduler = ExponentialLR(optimizer, gamma=0.9)

for epoch in range(20):
    for input, target in dataset:
        optimizer.zero_grad()
        output = model(input)
        loss = loss_fn(output, target)
        loss.backward()
        optimizer.step()
    scheduler.step()

Model.fc = nn.Identity() # method 1
Model.fc = nn.Sequential() # method 2
newmodel = torch.nn.Sequential(*(list(Model.children())[:-1])) # method 2

from utils.dist import init_distributed

init_distributed()

def init_distributed():
    dist_url = "env://" # default

    # only works with torch.distributed.launch // torch.run
    rank = int(os.environ["RANK"])
    world_size = int(os.environ['WORLD_SIZE'])
    local_rank = int(os.environ['LOCAL_RANK'])
    dist.init_process_group(
            backend="nccl",
            init_method=dist_url,
            world_size=world_size,
            rank=rank)
    
    torch.cuda.set_device(local_rank)
    dist.barrier()

from torch.utils.data.distributed import DistributedSampler
from utils.dist import get_rank

seed = args.seed + global_rank
torch.manual_seed(seed)
np.random.seed(seed)
random.seed(seed)

import torch.distributed as dist

def is_dist_avail_and_initialized():
    if not dist.is_available():
        return False
    if not dist.is_initialized():
        return False
    return True

def get_rank():
    if not is_dist_avail_and_initialized():
        return 0
    return dist.get_rank()

from torch.utils.data.distributed import DistributedSampler

train_data, valid_data = load_data(config.data.input_path, col_name=config.data.col_name,)
train_set, test_set = CrossDataset(train_data), CrossDataset(valid_data)



train_sampler = DistributedSampler(dataset=train_set, shuffle=True, rank=global_rank)
train_dataloader = DataLoader(train_set, batch_size=config.data.batch_size, sampler=train_sampler, num_workers=10, pin_memory=True)

test_sampler = DistributedSampler(dataset=test_set, shuffle=False, rank=global_rank)
test_dataloader = DataLoader(test_set, batch_size=config.data.batch_size, sampler=test_sampler, shuffle=False, num_workers=10, pin_memory=True)

if torch.cuda.is_available():  # for distributed parallel
    self.model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(self.model.cuda())
    local_rank = int(os.environ['LOCAL_RANK'])
    self.model = torch.nn.parallel.DistributedDataParallel(self.model, device_ids=[local_rank])

loader.sampler.set_epoch(epoch)  # for distributed parallel

if not is_dist_avail_and_initialized() or is_main_process():  # for distributed parallel
    save_model(self.model, base_dir, base_name)

def is_main_process():
    return get_rank() == 0

# 单击双卡运行
torchrun --nproc_per_node=2 train_bert.py # ...

git clone https://github.com/NVIDIA/apex
cd apex
pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --config-settings "--build-option=--cpp_ext" --config-settings "--build-option=--cuda_ext" ./

# Allow Amp to perform casts as required by the opt_level
model, optimizer = amp.initialize(model, optimizer, opt_level="O1")
self.model = nn.parallel.DistributedDataParallel(self.model, device_ids=[local_rank]) # 分布式计算，model parallel 需要在 apex.initialize 之后
...
# loss.backward() becomes:
with amp.scale_loss(loss, optimizer) as scaled_loss:
    scaled_loss.backward()
...

with torch.autocast(device_type=self.device, dtype=torch.float16, enabled=self.use_amp):
    loss = model.forward(x, y)

model.zero_grad()

with torch.autocast(device_type=self.device, dtype=torch.float16, enabled=self.use_amp):
  ...
  scaler.scale(loss).backward(retain_graph=True)

scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), self.grad_norm_clip)
scaler.step(optimizer)
scaler.update()
scheduler.step()
optimizer.zero_grad(set_to_none=True)

class LitGPT(pl.LightningModule):
    """ Use Pytorch Lightning wrapping to enable distributed parallel. """

    def __init__(self, vocab_size, n_embd, block_size, n_layer=8, n_head=8,
                 attn_pdrop=0.1, resid_pdrop=0.1, embd_pdrop=0.1,
                 learning_rate=3e-4, weight_decay=0.1, betas=(0.9, 0.95),  # used to config optimizer
                 warmup_tokens=375e6, final_tokens=260e9):
        super().__init__()
        self.save_hyperparameters()
        self.model = GPT(vocab_size, n_embd, block_size, n_layer=n_layer, n_head=n_head,
                         attn_pdrop=attn_pdrop, resid_pdrop=resid_pdrop, embd_pdrop=embd_pdrop)

    def forward(self, batch_x, targets=None):
        # Forward function that is run when visualizing the graph
        return self.model(batch_x, targets)

    def configure_optimizers(self):
        optimizer = self.model.configure_optimizers(learning_rate=self.hparams.learning_rate,
                                                    weight_decay=self.hparams.weight_decay,
                                                    betas=self.hparams.betas)

        t_mult = self.hparams.final_tokens - self.hparams.warmup_tokens
        scheduler = CosineAnnealingWarmRestarts(optimizer, int(self.hparams.warmup_tokens), T_mult=int(t_mult))

        return [optimizer], [scheduler]

    def _calculate_loss(self, batch, mode="train"):
        inp_x, labels = batch
        logits, loss = self.forward(inp_x, labels)
        acc = (logits.argmax(dim=-1) == labels).float().mean()

        self.log(f"{mode}_loss", loss)
        self.log(f"{mode}_acc", acc)

        return loss

    def training_step(self, batch, batch_idx):
        loss = self._calculate_loss(batch, mode="train")
        return loss

    @torch.no_grad()
    def validation_step(self, batch, batch_idx):
        _ = self._calculate_loss(batch, mode="val")

    @torch.no_grad()
    def test_step(self, batch, batch_idx):
        _ = self._calculate_loss(batch, mode="test")

def train(training_set: List[str], validation_set: List[str], output_dir,
          max_len=100, batch_size=64, n_layer=8, n_head=8, n_embd=512, attn_pdrop=0.1, resid_pdrop=0.1, embd_pdrop=0.1,
          n_epochs=10, lr=1e-3, weight_decay=0.1, betas=(0.9, 0.95), grad_clip_val=1.0, warmup_tokens=375e6,
          final_tokens=260e9, device='cuda', n_nodes=1, n_gpus=-1, ckpt_path=None, tf_train=True):

    prot_dm = ProtDataModule(training_set, validation_set, batch_size, max_len, num_workers=100)
    logger.info(f"Train dataloader length {len(prot_dm.train_dataloader().dataset)}, test dataloader length {len(prot_dm.val_dataloader().dataset)}")

    # Init trainer
    trainer = pl.Trainer(strategy=DDPStrategy(find_unused_parameters=False),
                         default_root_dir=output_dir,
                         callbacks=[ModelCheckpoint(save_weights_only=True, mode="max", monitor="train_acc"),
                                    TQDMProgressBar(refresh_rate=1), LearningRateMonitor("epoch")],  # to update
                         num_nodes=n_nodes, gpus=n_gpus if str(device).startswith("cuda") else 0, auto_select_gpus=True,
                         max_epochs=n_epochs, gradient_clip_val=grad_clip_val, log_every_n_steps=10)
    trainer.logger._default_hp_metric = None  # Optional logging argument that we don't need

    # Train model if no pretrained model exist
    if ckpt_path:
        print("Found pretrained model, loading...")
        model = LitGPT.load_from_checkpoint(ckpt_path)
    else:
        sd = AASeqDictionary()
        vocab_size = sd.get_char_num()
        block_size = max_len + 2  # add start & end

        model = LitGPT(vocab_size, n_embd, block_size, n_layer=n_layer, n_head=n_head,
                       attn_pdrop=attn_pdrop, resid_pdrop=resid_pdrop, embd_pdrop=embd_pdrop,
                       learning_rate=lr, weight_decay=weight_decay, betas=betas,
                       warmup_tokens=warmup_tokens, final_tokens=final_tokens)
    if tf_train:
        logger.info("check if model is training...")
        trainer.fit(model, prot_dm.train_dataloader())

    test_result = trainer.test(model, prot_dm.test_dataloader(), verbose=False)
    print(f"Test accuracy:  {(100.0 * test_result[0]['test_acc']):4.2f}%")

    return model.to(device)

class ProtDataModule(pl.LightningDataModule):
    def __init__(self, training_set: List[str], validation_set: List[str], batch_size=64, max_len=100, num_workers=100):
        super().__init__()
        self.batch_size = batch_size
        self.max_len = max_len
        self.num_workers = num_workers
        # load data
        train_seqs, _ = load_seqs_from_list(training_set, max_len=max_len, rm_duplicates=False)
        valid_seqs, _ = load_seqs_from_list(validation_set, max_len=max_len, rm_duplicates=False)
        logger.info(len(train_seqs))
        logger.info(f"Total GPU memory: {torch.cuda.get_device_properties(0).total_memory // 1024 ** 3} G")

        self.train_set = get_tensor_dataset(train_seqs)
        self.test_set = get_tensor_dataset(valid_seqs)
        logger.info(f"Train set length {len(self.train_set)}, test set length {len(self.test_set)}")

    def train_dataloader(self):
        return data.DataLoader(self.train_set, batch_size=self.batch_size, shuffle=True, drop_last=True, pin_memory=False,
                               num_workers=self.num_workers)

    def val_dataloader(self):
        return data.DataLoader(self.test_set, batch_size=self.batch_size, num_workers=self.num_workers, pin_memory=False)

    def test_dataloader(self):
        return data.DataLoader(self.test_set, batch_size=self.batch_size, num_workers=self.num_workers, pin_memory=False)

torch.load('model_dict.pkl')

torch.load('model_dict.pkl', map_location=lambda storage, loc: storage)

torch.load('model_dic.pkl', map_location=lambda storage, loc: storage.cuda)

#CPU->GPU1   模型是GPU1，预加载的训练参数却是CPU：
torch.load('model_dic.pkl', map_location=lambda storage, loc: storage.cuda(1))

from tqdm import tqdm

# Set tqdm
pbar = tqdm(enumerate(loader), total=len(loader)) if is_train else enumerate(loader)            

# Report progress
for epoch in range(n_epochs):
  for it, (x, y) in pbar:
    ...
    pbar.set_description(f"epoch {epoch + 1} iter {it}: train loss {loss.item():.5f}. lr {lr:e}")
  
# Result
# epoch 1 iter 11252: train loss 0.30460. lr 8.118256e-04:  57%|xxxx     | 11251/19893 

from torch.utils.tensorboard import SummaryWriter

# Set writer, usually in trainer
writer = SummaryWriter(self.config.output_dir + 'logs')          

# Report progress
for epoch in range(n_epochs):
  writer.add_scalar('loss', loss, epoch + 1)
  
# View results
tensorboard --logdir=./path/to/the/folder --port 8123 # 默认 30s 刷新一次

lsof /dev/nvidia1 | awk '{print $2}' | xargs -I {} kill -KILL {}
# show process using GPU card | show PID | kill process
# only need to kill the ones using the nvidia1

# Start
sudo docker compose ps

# Connect to Milvus
docker port milvus-standalone 19530/tcp

# Stop Milvus
sudo docker compose down

# Delete data after stop Milvus
sudo rm -rf  volumes

from pymilvus import connections, utility, FieldSchema, CollectionSchema, DataType, Collection

# Connect to Milvus server
connections.connect("default", host="localhost", port="19530")

# Creates a collection
fields = [
    FieldSchema(name="read_id", dtype=DataType.VARCHAR, max_length=40, is_primary=True, auto_id=False),
    FieldSchema(name="gene_id", dtype=DataType.VARCHAR, max_length=40, auto_id=False),
    FieldSchema(name="gene_name", dtype=DataType.VARCHAR, max_length=40, auto_id=False),
    FieldSchema(name="gene_type", dtype=DataType.VARCHAR, max_length=40, auto_id=False),
    FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=512, description="vector"),
]
schema = CollectionSchema(fields, "Signal annotation library")
collection = Collection("signal_annotation_library", schema)

# connect to collection
collection = Collection("signal_annotation_library")

# drop a collection
utility.drop_collection("signal_annotation_library")

# write to a collection
entities = [read_ids, gene_ids, gene_names, gene_types, sig_embds]

collection.insert(entities)
collection.flush()

# Deletes entities by their primary keys
expr = f"pk in [{entities[0][0]}, {entities[0][1]}]"
collection.delete(expr)

from pymilvus import Collection

# Builds indexes on the entities
index = {
    "index_type": "IVF_FLAT",
    "metric_type": "L2",
    "params": {"nlist": 512},
}
collection.create_index("embedding", index)

# Drop collection index
collection = Collection("signal_annotation_library")
collection.release()  # Release the collection.
collection.drop_index()

# Load the collection to memory
collection.load()

# Perform a vector similarity search
vectors_to_search = entities[-1][-2:]
search_params = {
    "metric_type": "L2",
    "params": {"nprobe": 10},
}

result = collection.search(vectors_to_search, "embeddings", search_params, limit=3, output_fields=["random"])