def __init__(self):
self.logger = settings.logger
torch.manual_seed(1)
torch.cuda.manual_seed_all(1)
torch.cuda.set_device(settings.GPU_ID)
if settings.DATASET == "WIKI":
self.train_dataset = datasets.WIKI(root=settings.DATA_DIR, train=True, transform=datasets.wiki_train_transform)
self.test_dataset = datasets.WIKI(root=settings.DATA_DIR, train=False, transform=datasets.wiki_test_transform)
self.database_dataset = datasets.WIKI(root=settings.DATA_DIR, train=True, transform=datasets.wiki_test_transform)
if settings.DATASET == "MIRFlickr":
self.train_dataset = datasets.MIRFlickr(train=True, transform=datasets.mir_train_transform)
self.test_dataset = datasets.MIRFlickr(train=False, database=False, transform=datasets.mir_test_transform)
self.database_dataset = datasets.MIRFlickr(train=False, database=True, transform=datasets.mir_test_transform)
if settings.DATASET == "NUSWIDE":
self.train_dataset = datasets.NUSWIDE(train=True, transform=datasets.nus_train_transform)
self.test_dataset = datasets.NUSWIDE(train=False, database=False, transform=datasets.nus_test_transform)
self.database_dataset = datasets.NUSWIDE(train=False, database=True, transform=datasets.nus_test_transform)
# Data Loader (Input Pipeline)
self.train_loader = torch.utils.data.DataLoader(dataset=self.train_dataset,
batch_size=settings.BATCH_SIZE,
shuffle=True,
num_workers=settings.NUM_WORKERS,
drop_last=True)
self.test_loader = torch.utils.data.DataLoader(dataset=self.test_dataset,
batch_size=settings.BATCH_SIZE,
shuffle=False,
num_workers=settings.NUM_WORKERS)
self.database_loader = torch.utils.data.DataLoader(dataset=self.database_dataset,
batch_size=settings.BATCH_SIZE,
shuffle=False,
num_workers=settings.NUM_WORKERS)
self.CodeNet_I = ImgNet(code_len=settings.CODE_LEN)
self.FeatNet_I = ImgNet(code_len=settings.CODE_LEN)
txt_feat_len = datasets.txt_feat_len
self.CodeNet_T = TxtNet(code_len=settings.CODE_LEN, txt_feat_len=txt_feat_len)
if settings.DATASET == "WIKI":
self.opt_I = torch.optim.SGD([{'params': self.CodeNet_I.fc_encode.parameters(), 'lr': settings.LR_IMG},
{'params': self.CodeNet_I.alexnet.classifier.parameters(), 'lr': settings.LR_IMG}],
momentum=settings.MOMENTUM, weight_decay=settings.WEIGHT_DECAY)
if settings.DATASET == "MIRFlickr" or settings.DATASET == "NUSWIDE":
self.opt_I = torch.optim.SGD(self.CodeNet_I.parameters(), lr=settings.LR_IMG, momentum=settings.MOMENTUM, weight_decay=settings.WEIGHT_DECAY)
self.opt_T = torch.optim.SGD(self.CodeNet_T.parameters(), lr=settings.LR_TXT, momentum=settings.MOMENTUM, weight_decay=settings.WEIGHT_DECAY)
def __init__(self, log, config):
self.logger = log
self.config = config
torch.manual_seed(1)
torch.cuda.manual_seed_all(1)
torch.cuda.set_device(self.config.GPU_ID)
if self.config.DATASET == "MIRFlickr":
self.train_dataset = datasets.MIRFlickr(train=True, transform=datasets.mir_train_transform)
self.test_dataset = datasets.MIRFlickr(train=False, database=False, transform=datasets.mir_test_transform)
self.database_dataset = datasets.MIRFlickr(train=False, database=True,
transform=datasets.mir_test_transform)
if self.config.DATASET == "NUSWIDE":
self.train_dataset = datasets.NUSWIDE(train=True, transform=datasets.nus_train_transform)
self.test_dataset = datasets.NUSWIDE(train=False, database=False, transform=datasets.nus_test_transform)
self.database_dataset = datasets.NUSWIDE(train=False, database=True, transform=datasets.nus_test_transform)
# Data Loader (Input Pipeline)
self.train_loader = torch.utils.data.DataLoader(dataset=self.train_dataset,
batch_size=self.config.BATCH_SIZE,
shuffle=True,
num_workers=self.config.NUM_WORKERS,
drop_last=True)
self.test_loader = torch.utils.data.DataLoader(dataset=self.test_dataset,
batch_size=self.config.BATCH_SIZE,
shuffle=False,
num_workers=self.config.NUM_WORKERS)
self.database_loader = torch.utils.data.DataLoader(dataset=self.database_dataset,
batch_size=self.config.BATCH_SIZE,
shuffle=False,
num_workers=self.config.NUM_WORKERS)
self.ImgNet = ImgNet(code_len=self.config.HASH_BIT)
txt_feat_len = datasets.txt_feat_len
self.TxtNet = TxtNet(code_len=self.config.HASH_BIT, txt_feat_len=txt_feat_len)
self.opt_I = torch.optim.SGD(self.ImgNet.parameters(), lr=self.config.LR_IMG, momentum=self.config.MOMENTUM,
weight_decay=self.config.WEIGHT_DECAY)
self.opt_T = torch.optim.SGD(self.TxtNet.parameters(), lr=self.config.LR_TXT, momentum=self.config.MOMENTUM,
weight_decay=self.config.WEIGHT_DECAY)
self.best_it = 0
self.best_ti = 0
def __init__(self, log, cfg):
self.since = time.time()
self.logger = log
self.cfg = cfg
self.epoch_loss = 0.
self.path = "_".join([
self.cfg.MODEL,
str(self.cfg.HASH_BIT), self.cfg.DATASET, self.cfg.TAG
])
torch.manual_seed(1)
torch.cuda.manual_seed_all(1)
torch.cuda.set_device(self.cfg.GPU_ID)
if self.cfg.DATASET == "MIRFlickr":
self.train_dataset = datasets.MIRFlickr(
train=True, transform=datasets.mir_train_transform)
self.test_dataset = datasets.MIRFlickr(
train=False,
database=False,
transform=datasets.mir_test_transform)
self.database_dataset = datasets.MIRFlickr(
train=False,
database=True,
transform=datasets.mir_test_transform)
if self.cfg.DATASET == "NUSWIDE":
self.train_dataset = datasets.NUSWIDE(
train=True, transform=datasets.nus_train_transform)
self.test_dataset = datasets.NUSWIDE(
train=False,
database=False,
transform=datasets.nus_test_transform)
self.database_dataset = datasets.NUSWIDE(
train=False,
database=True,
transform=datasets.nus_test_transform)
if self.cfg.DATASET == "UCM" or "RSICD":
da = self.cfg.DATA_AMOUNT
self.train_dataset = datasets.UCM2(type='train', data_amount=da)
self.test_dataset = datasets.UCM2(type='query', data_amount=da)
self.database_dataset = datasets.UCM2(type='db', data_amount=da)
# Data Loader (Input Pipeline)
self.train_loader = torch.utils.data.DataLoader(
dataset=self.train_dataset,
batch_size=self.cfg.BATCH_SIZE,
shuffle=True,
num_workers=self.cfg.NUM_WORKERS,
drop_last=True)
self.test_loader = torch.utils.data.DataLoader(
dataset=self.test_dataset,
batch_size=self.cfg.BATCH_SIZE,
shuffle=False,
num_workers=self.cfg.NUM_WORKERS)
self.database_loader = torch.utils.data.DataLoader(
dataset=self.database_dataset,
batch_size=self.cfg.BATCH_SIZE,
shuffle=False,
num_workers=self.cfg.NUM_WORKERS)
if self.cfg.DATASET == "UCM" or "RSICD":
txt_feat_len = datasets.txt_feat_len
img_feat_len = datasets.img_feat_len
self.ImgNet = ImgNetRS(self.cfg.HASH_BIT, img_feat_len, 512)
self.TxtNet = TxtNetRS(self.cfg.HASH_BIT, txt_feat_len, 512)
else:
txt_feat_len = datasets.txt_feat_len
self.ImgNet = ImgNet(code_len=self.cfg.HASH_BIT)
self.TxtNet = TxtNet(code_len=self.cfg.HASH_BIT,
txt_feat_len=txt_feat_len)
self.opt_I = torch.optim.SGD(self.ImgNet.parameters(),
lr=self.cfg.LR_IMG,
momentum=self.cfg.MOMENTUM,
weight_decay=self.cfg.WEIGHT_DECAY)
self.opt_T = torch.optim.SGD(self.TxtNet.parameters(),
lr=self.cfg.LR_TXT,
momentum=self.cfg.MOMENTUM,
weight_decay=self.cfg.WEIGHT_DECAY)
self.best_it = 0
self.best_ti = 0
self.best_ii = 0
self.best_tt = 0