From a22408c1ecb1d04ee80a9bc288537acc7805364c Mon Sep 17 00:00:00 2001 From: leewlving Date: Sat, 8 Jun 2024 14:45:18 +0800 Subject: [PATCH] update --- train/hash_train.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/train/hash_train.py b/train/hash_train.py index c3c5bbc..f0e9b42 100644 --- a/train/hash_train.py +++ b/train/hash_train.py @@ -17,6 +17,8 @@ from dataset.dataloader import dataloader import open_clip # from transformers import BertModel +device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") + def clamp(delta, clean_imgs): clamp_imgs = (delta.data + clean_imgs.data).clamp(0, 1) @@ -34,6 +36,7 @@ class Trainer(TrainBase): image_representation, text_representation=self.generate_mapping() self.image_representation=image_representation self.text_representation=text_representation + self.device=rank # self.run() def _init_model(self): @@ -74,7 +77,7 @@ class Trainer(TrainBase): # if self.args.pretrained != "" and os.path.exists(self.args.pretrained): # self.logger.info("load pretrained model.") # self.model.load_state_dict(torch.load(self.args.pretrained, map_location=f"cuda:{self.rank}")) - model_clip, _, preprocess = open_clip.create_model_and_transforms('ViT-B-16', device=self.device) + model_clip, _, preprocess = open_clip.create_model_and_transforms('ViT-B-16', device=device) self.model= model_clip self.model.eval() self.model.float() @@ -132,11 +135,9 @@ class Trainer(TrainBase): def generate_mapping(self): text_train=[] label_train=[] - # image_train=[] - # self.change_state(mode="valid") for image, text, label, index in self.train_loader: # image=image.to(self.device, non_blocking=True) - text=text.to(self.device, non_blocking=True) + text=text.to(device, non_blocking=True) temp_text=self.model.encode_text(text) # temp_image=self.model.encode_image(image) # image_train.append(temp_image.cpu().detach().numpy())