From c479c07a2cd660c49eef8559db56476605c140a4 Mon Sep 17 00:00:00 2001 From: leewlving Date: Sat, 8 Jun 2024 16:59:39 +0800 Subject: [PATCH] new update --- model/GAN.py | 310 ----------------------------------------- model/hash_model.py | 161 --------------------- model/spectral_norm.py | 89 ------------ train/hash_train.py | 20 +-- utils/get_args.py | 2 +- 5 files changed, 12 insertions(+), 570 deletions(-) delete mode 100644 model/GAN.py delete mode 100644 model/hash_model.py delete mode 100644 model/spectral_norm.py diff --git a/model/GAN.py b/model/GAN.py deleted file mode 100644 index f7282ee..0000000 --- a/model/GAN.py +++ /dev/null @@ -1,310 +0,0 @@ -import torch -import torch.nn as nn -import torch.nn.functional as F -from torch.optim import lr_scheduler - -from model.spectral_norm import spectral_norm as SpectralNorm - - -class PrototypeNet(nn.Module): - def __init__(self, bit, num_classes): - super(PrototypeNet, self).__init__() - - self.feature = nn.Sequential(nn.Linear(num_classes, 4096), - nn.ReLU(True), nn.Linear(4096, 512)) - self.hashing = nn.Sequential(nn.Linear(512, bit), nn.Tanh()) - self.classifier = nn.Sequential(nn.Linear(512, num_classes), - nn.Sigmoid()) - - def forward(self, label): - f = self.feature(label) - h = self.hashing(f) - c = self.classifier(f) - return f, h, c - - -class Discriminator(nn.Module): - """ - Discriminator network with PatchGAN. - Reference: https://github.com/yunjey/stargan/blob/master/model.py - """ - def __init__(self, num_classes, image_size=224, conv_dim=64, repeat_num=5): - super(Discriminator, self).__init__() - layers = [] - layers.append(SpectralNorm(nn.Conv2d(3, conv_dim, kernel_size=4, stride=2, padding=1))) - layers.append(nn.LeakyReLU(0.01)) - - curr_dim = conv_dim - for i in range(1, repeat_num): - layers.append(SpectralNorm(nn.Conv2d(curr_dim, curr_dim*2, kernel_size=4, stride=2, padding=1))) - layers.append(nn.LeakyReLU(0.01)) - curr_dim = curr_dim * 2 - - kernel_size = int(image_size / (2**repeat_num)) - self.main = nn.Sequential(*layers) - self.fc = nn.Conv2d(curr_dim, num_classes + 1, kernel_size=kernel_size, bias=False) - - def forward(self, x): - h = self.main(x) - out = self.fc(h) - return out.squeeze() - - -class Generator(nn.Module): - """Generator: Encoder-Decoder Architecture. - Reference: https://github.com/yunjey/stargan/blob/master/model.py - """ - def __init__(self): - super(Generator, self).__init__() - - # Label Encoder - self.label_encoder = LabelEncoder() - - # Image Encoder - curr_dim = 64 - image_encoder = [ - nn.Conv2d(6, curr_dim, kernel_size=7, stride=1, padding=3, bias=True), - nn.InstanceNorm2d(curr_dim), - nn.ReLU(inplace=True) - ] - # Down Sampling - for i in range(2): - image_encoder += [ - nn.Conv2d(curr_dim, - curr_dim * 2, - kernel_size=4, - stride=2, - padding=1, - bias=True), - nn.InstanceNorm2d(curr_dim * 2), - nn.ReLU(inplace=True) - ] - curr_dim = curr_dim * 2 - # Bottleneck - for i in range(3): - image_encoder += [ - ResidualBlock(dim_in=curr_dim, dim_out=curr_dim, net_mode='t') - ] - self.image_encoder = nn.Sequential(*image_encoder) - - # Decoder - decoder = [] - # Bottleneck - for i in range(3): - decoder += [ - ResidualBlock(dim_in=curr_dim, dim_out=curr_dim, net_mode='t') - ] - # Up Sampling - for i in range(2): - decoder += [ - nn.ConvTranspose2d(curr_dim, - curr_dim // 2, - kernel_size=4, - stride=2, - padding=1, - bias=False), - nn.InstanceNorm2d(curr_dim // 2), - nn.ReLU(inplace=True) - ] - curr_dim = curr_dim // 2 - self.residual = nn.Sequential( - # nn.Conv2d(curr_dim + 3, - # curr_dim, - # kernel_size=3, - # stride=1, - # padding=1, - # bias=False), - # nn.InstanceNorm2d(curr_dim // 2, affine=False), - # nn.ReLU(inplace=True), - nn.Conv2d(curr_dim + 3, - 3, - kernel_size=3, - stride=1, - padding=1, - bias=False), nn.Tanh()) - self.decoder = nn.Sequential(*decoder) - - def forward(self, x, label_feature): - mixed_feature = self.label_encoder(x, label_feature) - encode = self.image_encoder(mixed_feature) - decode = self.decoder(encode) - decode_x = torch.cat([decode, x], dim=1) - adv_x = self.residual(decode_x) - return adv_x, mixed_feature - - -class LabelEncoder(nn.Module): - def __init__(self, nf=128): - super(LabelEncoder, self).__init__() - self.nf = nf - curr_dim = nf - self.size = 14 - - self.fc = nn.Sequential( - # nn.Linear(512, 512), nn.ReLU(True), - nn.Linear(512, curr_dim * self.size * self.size), nn.ReLU(True)) - - transform = [] - for i in range(4): - transform += [ - nn.ConvTranspose2d(curr_dim, - curr_dim // 2, - kernel_size=4, - stride=2, - padding=1, - bias=False), - # nn.Upsample(scale_factor=(2, 2)), - # nn.Conv2d(curr_dim, curr_dim//2, kernel_size=3, padding=1, bias=False), - nn.InstanceNorm2d(curr_dim // 2, affine=False), - nn.ReLU(inplace=True) - ] - curr_dim = curr_dim // 2 - - transform += [ - nn.Conv2d(curr_dim, - 3, - kernel_size=3, - stride=1, - padding=1, - bias=False) - ] - self.transform = nn.Sequential(*transform) - - def forward(self, image, label_feature): - label_feature = self.fc(label_feature) - label_feature = label_feature.view(label_feature.size(0), self.nf, self.size, self.size) - label_feature = self.transform(label_feature) - - # mixed_feature = label_feature + image - mixed_feature = torch.cat((label_feature, image), dim=1) - return mixed_feature - - -class ResidualBlock(nn.Module): - """Residual Block.""" - def __init__(self, dim_in, dim_out, net_mode=None): - if net_mode == 'p' or (net_mode is None): - use_affine = True - elif net_mode == 't': - use_affine = False - super(ResidualBlock, self).__init__() - self.main = nn.Sequential( - nn.Conv2d(dim_in, - dim_out, - kernel_size=3, - stride=1, - padding=1, - bias=False), nn.InstanceNorm2d(dim_out, - affine=use_affine), - nn.ReLU(inplace=True), - nn.Conv2d(dim_out, - dim_out, - kernel_size=3, - stride=1, - padding=1, - bias=False), nn.InstanceNorm2d(dim_out, - affine=use_affine)) - - def forward(self, x): - return x + self.main(x) - - -class GANLoss(nn.Module): - """Define different GAN objectives. - The GANLoss class abstracts away the need to create the target label tensor - that has the same size as the input. - """ - def __init__(self, gan_mode, target_real_label=0.0, target_fake_label=1.0): - """ Initialize the GANLoss class. - Parameters: - gan_mode (str) - - the type of GAN objective. It currently supports vanilla, lsgan, and wgangp. - target_real_label (bool) - - label for a real image - target_fake_label (bool) - - label of a fake image - Note: Do not use sigmoid as the last layer of Discriminator. - LSGAN needs no sigmoid. vanilla GANs will handle it with BCEWithLogitsLoss. - """ - super(GANLoss, self).__init__() - self.register_buffer('real_label', torch.tensor(target_real_label)) - self.register_buffer('fake_label', torch.tensor(target_fake_label)) - self.gan_mode = gan_mode - if gan_mode == 'lsgan': - self.loss = nn.MSELoss() - elif gan_mode == 'vanilla': - self.loss = nn.BCEWithLogitsLoss() - elif gan_mode in ['wgangp']: - self.loss = None - else: - raise NotImplementedError('gan mode %s not implemented' % gan_mode) - - def get_target_tensor(self, label, target_is_real): - """Create label tensors with the same size as the input. - Parameters: - prediction (tensor) - - tpyically the prediction from a discriminator - target_is_real (bool) - - if the ground truth label is for real images or fake images - Returns: - A label tensor filled with ground truth label, and with the size of the input - """ - if target_is_real: - real_label = self.real_label.expand(label.size(0), 1) - target_tensor = torch.cat([label, real_label], dim=-1) - else: - fake_label = self.fake_label.expand(label.size(0), 1) - target_tensor = torch.cat([label, fake_label], dim=-1) - return target_tensor - - def __call__(self, prediction, label, target_is_real): - """Calculate loss given Discriminator's output and grount truth labels. - Parameters: - prediction (tensor) - - tpyically the prediction output from a discriminator - target_is_real (bool) - - if the ground truth label is for real images or fake images - Returns: - the calculated loss. - """ - if self.gan_mode in ['lsgan', 'vanilla']: - target_tensor = self.get_target_tensor(label, target_is_real) - loss = self.loss(prediction, target_tensor) - elif self.gan_mode == 'wgangp': - if target_is_real: - loss = -prediction.mean() - else: - loss = prediction.mean() - return loss - - -def get_scheduler(optimizer, opt): - """Return a learning rate scheduler - Parameters: - optimizer -- the optimizer of the network - opt (option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions. - opt.lr_policy is the name of learning rate policy: linear | step | plateau | cosine - For 'linear', we keep the same learning rate for the first epochs - and linearly decay the rate to zero over the next epochs. - For other schedulers (step, plateau, and cosine), we use the default PyTorch schedulers. - See https://pytorch.org/docs/stable/optim.html for more details. - """ - if opt.lr_policy == 'linear': - - def lambda_rule(epoch): - lr_l = 1.0 - max(0, epoch + opt.epoch_count - - opt.n_epochs) / float(opt.n_epochs_decay + 1) - return lr_l - - scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda_rule) - elif opt.lr_policy == 'step': - scheduler = lr_scheduler.StepLR(optimizer, - step_size=opt.lr_decay_iters, - gamma=0.1) - elif opt.lr_policy == 'plateau': - scheduler = lr_scheduler.ReduceLROnPlateau(optimizer, - mode='min', - factor=0.2, - threshold=0.01, - patience=5) - elif opt.lr_policy == 'cosine': - scheduler = lr_scheduler.CosineAnnealingLR(optimizer, - T_max=opt.n_epochs, - eta_min=0) - else: - return NotImplementedError( - 'learning rate policy [%s] is not implemented', opt.lr_policy) - return scheduler \ No newline at end of file diff --git a/model/hash_model.py b/model/hash_model.py deleted file mode 100644 index 2d9c851..0000000 --- a/model/hash_model.py +++ /dev/null @@ -1,161 +0,0 @@ -import os -import torch -import logging -import torch.nn as nn -import numpy as np -from typing import Union - -from model.model import build_model -from utils import get_logger, get_summary_writer - -def weights_init_kaiming(m): - classname = m.__class__.__name__ - if classname.find('Linear') != -1: - nn.init.kaiming_uniform_(m.weight, mode='fan_out') - nn.init.constant_(m.bias, 0.0) - elif classname.find('Conv') != -1: - nn.init.kaiming_normal_(m.weight, a=0, mode='fan_in') - if m.bias is not None: - nn.init.constant_(m.bias, 0.0) - elif classname.find('BatchNorm') != -1: - if m.affine: - nn.init.constant_(m.weight, 1.0) - nn.init.constant_(m.bias, 0.0) - - -class LinearHash(nn.Module): - - def __init__(self, inputDim=2048, outputDim=64): - super(LinearHash, self).__init__() - self.fc = nn.Linear(inputDim, outputDim) - self.fc.apply(weights_init_kaiming) - self.drop_out = nn.Dropout(p=0.2) - - def forward(self, data): - result = self.fc(data) - return torch.tanh(self.drop_out(result)) - - -class HashLayer(nn.Module): - - LINEAR_EMBED = 128 - SIGMOID_ALPH = 10 - def __init__(self, inputDim=2048, outputDim=64): - - super(HashLayer, self).__init__() - self.fc = nn.Linear(inputDim, self.LINEAR_EMBED) - self.fc.apply(weights_init_kaiming) - self.hash_list = nn.ModuleList([nn.Linear(self.LINEAR_EMBED, 2) for _ in range(outputDim)]) - for item in self.hash_list: - item.apply(weights_init_kaiming) - - def forward(self, data): - - embed = self.fc(data) - embed = torch.relu(embed) - - softmax_list = [torch.softmax(item(embed), dim=-1) for item in self.hash_list] - - return softmax_list - -class HashLayer_easy_logic(nn.Module): - - LINEAR_EMBED = 128 - SIGMOID_ALPH = 10 - def __init__(self, inputDim=2048, outputDim=64): - - super(HashLayer, self).__init__() - self.bit = outputDim - self.fc = nn.Linear(inputDim, outputDim * 2) - self.fc.apply(weights_init_kaiming) - for item in self.hash_list: - item.apply(weights_init_kaiming) - - def forward(self, data): - - embed = self.fc(data) - softmax_list = embed.view(embed.shape[0], self.bit, 2) - - softmax_list = torch.softmax(softmax_list, dim=-1) - - return softmax_list - - -class DCMHT(nn.Module): - - def __init__(self, - outputDim=64, - clipPath="./ViT-B-32.pt", - writer=None, - saveDir="./result/log", - logger: logging.Logger=None, - is_train=True, - linear=False): - super(DCMHT, self).__init__() - os.makedirs(saveDir, exist_ok=True) - self.logger = logger if logger is not None else get_logger(os.path.join(saveDir, "train.log" if is_train else "test.log")) - self.writer = writer if writer is not None and is_train else get_summary_writer(os.path.join(saveDir, "tensorboard")) - embedDim, self.clip = self.load_clip(clipPath) - # if is_train: - # self.clip.eval() - # print("start freezen") - # self.freezen() - self.image_hash = LinearHash(inputDim=embedDim, outputDim=outputDim) if linear else HashLayer(inputDim=embedDim, outputDim=outputDim) - self.text_hash = LinearHash(inputDim=embedDim, outputDim=outputDim) if linear else HashLayer(inputDim=embedDim, outputDim=outputDim) - # print(self.image_hash) - # print(self.text_hash) - - def freezen(self): - for name, param in self.clip.named_parameters(): - # print(name) - if name.find("ln_final.") == 0 or name.find("text_projection") == 0 or name.find("logit_scale") == 0 \ - or name.find("visual.ln_post.") == 0 or name.find("visual.proj") == 0: - # print("1") - continue - elif name.find("visual.transformer.resblocks.") == 0 or name.find("transformer.resblocks.") == 0: - layer_num = int(name.split(".resblocks.")[1].split(".")[0]) - if layer_num >= 12: - # print("2") - continue - if name.find("conv2.") == 0: - # print("3") - continue - else: - # paramenters which < freeze_layer_num will be freezed - param.requires_grad = False - - def load_clip(self, clipPath: str) -> tuple: - try: - model = torch.jit.load(clipPath, map_location="cpu").eval() - state_dict = model.state_dict() - except RuntimeError: - state_dict = torch.load(clipPath, map_location="cpu") - - return state_dict["text_projection"].shape[1], build_model(state_dict) - - def encode_image(self, image): - - image_embed = self.clip.encode_image(image) - image_embed = self.image_hash(image_embed) - - return image_embed - - def eval(self): - self.image_hash.eval() - self.text_hash.eval() - # self.clip.eval() - - def train(self): - self.image_hash.train() - self.text_hash.train() - - def encode_text(self, text): - - text_embed = self.clip.encode_text(text) - text_embed = self.text_hash(text_embed) - - return text_embed - - def forward(self, image, text): - return self.encode_image(image), self.encode_text(text) - diff --git a/model/spectral_norm.py b/model/spectral_norm.py deleted file mode 100644 index 1f219b6..0000000 --- a/model/spectral_norm.py +++ /dev/null @@ -1,89 +0,0 @@ -import torch -from torch.nn import Parameter - - -def l2normalize(v, eps=1e-12): - return v / (v.norm() + eps) - - -class SpectralNorm(object): - def __init__(self): - self.name = "weight" - #print(self.name) - self.power_iterations = 1 - - def compute_weight(self, module): - u = getattr(module, self.name + "_u") - v = getattr(module, self.name + "_v") - w = getattr(module, self.name + "_bar") - - height = w.data.shape[0] - for _ in range(self.power_iterations): - v.data = l2normalize( - torch.mv(torch.t(w.view(height, -1).data), u.data)) - u.data = l2normalize(torch.mv(w.view(height, -1).data, v.data)) - # sigma = torch.dot(u.data, torch.mv(w.view(height,-1).data, v.data)) - sigma = u.dot(w.view(height, -1).mv(v)) - return w / sigma.expand_as(w) - - @staticmethod - def apply(module): - name = "weight" - fn = SpectralNorm() - - try: - u = getattr(module, name + "_u") - v = getattr(module, name + "_v") - w = getattr(module, name + "_bar") - except AttributeError: - w = getattr(module, name) - height = w.data.shape[0] - width = w.view(height, -1).data.shape[1] - u = Parameter(w.data.new(height).normal_(0, 1), - requires_grad=False) - v = Parameter(w.data.new(width).normal_(0, 1), requires_grad=False) - w_bar = Parameter(w.data) - - #del module._parameters[name] - - module.register_parameter(name + "_u", u) - module.register_parameter(name + "_v", v) - module.register_parameter(name + "_bar", w_bar) - - # remove w from parameter list - del module._parameters[name] - - setattr(module, name, fn.compute_weight(module)) - - # recompute weight before every forward() - module.register_forward_pre_hook(fn) - - return fn - - def remove(self, module): - weight = self.compute_weight(module) - delattr(module, self.name) - del module._parameters[self.name + '_u'] - del module._parameters[self.name + '_v'] - del module._parameters[self.name + '_bar'] - module.register_parameter(self.name, Parameter(weight.data)) - - def __call__(self, module, inputs): - setattr(module, self.name, self.compute_weight(module)) - - -def spectral_norm(module): - SpectralNorm.apply(module) - return module - - -def remove_spectral_norm(module): - name = 'weight' - for k, hook in module._forward_pre_hooks.items(): - if isinstance(hook, SpectralNorm) and hook.name == name: - hook.remove(module) - del module._forward_pre_hooks[k] - return module - - raise ValueError("spectral_norm of '{}' not found in {}".format( - name, module)) \ No newline at end of file diff --git a/train/hash_train.py b/train/hash_train.py index 972fbfa..48f1332 100644 --- a/train/hash_train.py +++ b/train/hash_train.py @@ -191,18 +191,20 @@ class Trainer(TrainBase): def get_code(self, data_loader, length: int): - img_buffer = [] - text_buffer = [] + img_buffer = torch.empty(length, self.args.output_dim, dtype=torch.float).to(self.rank) + text_buffer = torch.empty(length, self.args.output_dim, dtype=torch.float).to(self.rank) for image, text, label, index in tqdm(data_loader): - image = image.to(self.rank, non_blocking=True) - text = text.to(self.rank, non_blocking=True) + image = image.to(self.device, non_blocking=True) + text = text.to(self.device, non_blocking=True) index = index.numpy() - image_hash=self.model.encode_image(image) - # text_feat=self.bert(text)[0] - text_hash=self.model.encode_text(text) - img_buffer[index, :] = image_hash.data - text_buffer[index, :] = text_hash.data + with torch.no_grad(): + image_feature = self.model.encode_image(image) + text_features = self.model.encode_text(text) + image_feature /= image_feature.norm(dim=-1, keepdim=True) + text_features /= text_features.norm(dim=-1, keepdim=True) + img_buffer[index, :] = image_feature.detach() + text_buffer[index, :] = text_features.detach() return img_buffer, text_buffer# img_buffer.to(self.rank), text_buffer.to(self.rank) diff --git a/utils/get_args.py b/utils/get_args.py index d8513fe..7216497 100644 --- a/utils/get_args.py +++ b/utils/get_args.py @@ -19,7 +19,7 @@ def get_args(): # parser.add_argument("--test-caption-file", type=str, default="./data/test/captions.mat") # parser.add_argument("--test-label-file", type=str, default="./data/test/label.mat") parser.add_argument("--txt-dim", type=int, default=1024) - parser.add_argument("--output-dim", type=int, default=64) + parser.add_argument("--output-dim", type=int, default=512) parser.add_argument("--epochs", type=int, default=100) parser.add_argument("--max-words", type=int, default=77) parser.add_argument("--resolution", type=int, default=224)