liwenyun
/
advclip
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

update

This commit is contained in:
leewlving 2024-06-08 14:45:18 +08:00
parent e79be260b4
commit a22408c1ec
1 changed files with 5 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
train/hash_train.py
View File

@ -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())