def load_model(checkpoints_folder,device):
model =ResNetSimCLR(**config['model'])
model.eval()
state_dict = torch.load(os.path.join(checkpoints_folder, 'model.pth'), map_location=torch.device('cpu'))
model.load_state_dict(state_dict)
model = model.to(device)
return model
def encode(save_root, model_file, data_folder, model_name='ca', dataset_name='celeba', batch_size=64, device='cuda:0', out_dim=256):
os.makedirs(save_root, exist_ok=True)
os.makedirs(data_folder, exist_ok=True)
if dataset_name == 'celeba':
train_loader = DataLoader(datasets.CelebA(data_folder, split='train', download=True, transform=transforms.ToTensor()),
batch_size=batch_size, shuffle=False)
valid_loader = DataLoader(datasets.CelebA(data_folder, split='valid', download=True, transform=transforms.ToTensor()),
batch_size=batch_size, shuffle=False)
elif dataset_name == 'stanfordCars':
t = transforms.Compose([
transforms.Resize(512),
transforms.CenterCrop(512),
transforms.ToTensor(),
transforms.Lambda(lambda x: x.repeat(3,1,1) if x.shape[0] == 1 else x)
])
train_data_dir = os.path.join(data_folder, 'cars_train/')
train_annos = os.path.join(data_folder, 'devkit/cars_train_annos.mat')
train_loader = DataLoader(CarsDataset(train_annos, train_data_dir, t), batch_size=batch_size, shuffle=False)
valid_data_dir = os.path.join(data_folder, 'cars_test/')
valid_annos = os.path.join(data_folder, 'devkit/cars_test_annos_withlabels.mat')
valid_loader = DataLoader(CarsDataset(valid_annos, valid_data_dir, t), batch_size=batch_size, shuffle=False)
elif dataset_name == 'compCars':
t = transforms.Compose([
transforms.Resize(512),
transforms.CenterCrop(512),
transforms.ToTensor()
])
train_loader = DataLoader(CompCars(data_folder, True, t), batch_size=batch_size, shuffle=False)
valid_loader = DataLoader(CompCars(data_folder, False, t), batch_size=batch_size, shuffle=False)
model = ResNetSimCLR('resnet50', out_dim)
model.load_state_dict(torch.load(model_file, map_location=device))
model = model.to(device)
model.eval()
print('Starting on training data')
train_encodings = []
for x, _ in train_loader:
x = x.to(device)
h, _ = model(x)
train_encodings.append(h.cpu().detach())
torch.save(torch.cat(train_encodings, dim=0), os.path.join(save_root, f'{dataset_name}-{model_name}model-train_encodings.pt'))
print('Starting on validation data')
valid_encodings = []
for x, _ in valid_loader:
x = x.to(device)
h, _ = model(x)
if len(h.shape) == 1:
h = h.unsqueeze(0)
valid_encodings.append(h.cpu().detach())
torch.save(torch.cat(valid_encodings, dim=0), os.path.join(save_root, f'{dataset_name}-{model_name}model-valid_encodings.pt'))
def run(config):
model = ResNetSimCLR('resnet50', config.out_dim)
model.load_state_dict(
torch.load(config.model_file, map_location=config.device))
model = model.to(config.device)
clf = nn.Linear(2048, 196)
full = FineTuner(model, clf)
optim = torch.optim.Adam(list(model.parameters()) + list(clf.parameters()))
objective = nn.CrossEntropyLoss()
t = T.Compose([
T.Resize(512),
T.CenterCrop(512),
T.ToTensor(),
T.Lambda(lambda x: x.repeat(3, 1, 1) if x.shape[0] == 1 else x)
])
train_data_dir = os.path.join(config.data_root, 'cars_train/')
train_annos = os.path.join(config.data_root, 'devkit/cars_train_annos.mat')
train_dataset = StanfordCarsMini(train_annos, train_data_dir, t)
train_loader = DataLoader(train_dataset,
batch_size=config.batch_size,
shuffle=True)
valid_data_dir = os.path.join(config.data_root, 'cars_test/')
valid_annos = os.path.join(config.data_root,
'devkit/cars_test_annos_withlabels.mat')
valid_dataset = CarsDataset(valid_annos, valid_data_dir, t)
valid_loader = DataLoader(valid_dataset, batch_size=config.batch_size)
solver = CESolver(full,
train_loader,
valid_loader,
config.save_root,
name=config.name,
device=config.device)
solver.train(config.num_epochs)
def evaluation(checkpoints_folder, config, device):
model = ResNetSimCLR(**config['model'])
model.eval()
model.load_state_dict(
torch.load(os.path.join(checkpoints_folder, 'model.pth')))
model = model.to(device)
train_set = torchvision.datasets.CIFAR10(
root='../data/CIFAR10',
transform=transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
]),
train=True,
download=True)
test_set = torchvision.datasets.CIFAR10(root='../data/CIFAR10',
transform=transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(
mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
]),
train=False,
download=True)
# num_train = len(train_dataset)
# indices = list(range(num_train))
# np.random.shuffle(indices)
#
# split = int(np.floor(0.05 * num_train))
# train_idx, test_idx = indices[split:], indices[:split]
#
# # define samplers for obtaining training and validation batches
# train_sampler = torch.utils.data.sampler.SubsetRandomSampler(train_idx)
# test_sampler = torch.utils.data.sampler.SubsetRandomSampler(test_idx) # ?????sampler????????????
# train_loader = torch.utils.data.DataLoader(train_set, batch_size=config['batch_size'], drop_last=True, shuffle=True)
#
# test_loader = torch.utils.data.DataLoader(test_set, batch_size=config['batch_size'], drop_last=True, shuffle=True)
train_loader = torch.utils.data.DataLoader(train_set,
batch_size=48,
drop_last=True,
shuffle=True)
test_loader = torch.utils.data.DataLoader(test_set,
batch_size=48,
drop_last=True,
shuffle=True)
X_train_feature = []
label_train = []
for data in train_loader:
x, y = data
x = x.to(device)
features, _ = model(x)
X_train_feature.extend(features.cpu().detach().numpy())
label_train.extend(y.cpu().detach().numpy())
X_train_feature = np.array(X_train_feature)
label_train = np.array(label_train)
X_test_feature = []
label_test = []
for data in test_loader:
x, y = data
x = x.to(device)
features, _ = model(x)
X_test_feature.extend(features.cpu().detach().numpy())
label_test.extend(y.cpu().detach().numpy())
X_test_feature = np.array(X_test_feature)
label_test = np.array(label_test)
scaler = preprocessing.StandardScaler()
print('ok')
scaler.fit(X_train_feature)
# print(X_test_feature.shape)
# print(y_test.shape)
linear_model_eval(scaler.transform(X_train_feature), label_train,
scaler.transform(X_test_feature), label_test)
def eval_frozen(self):
train_loader, val_loader, num_classes = self.dataset.get_dataset_eval()
model = ResNetSimCLR(**self.config["model"]).to(self.device)
model = self._load_pre_trained_weights(model)
model.to(self.device)
model.eval()
lineal_classifier = nn.Linear(model.l1.in_features, num_classes)
lineal_classifier.to(self.device)
optimizer = torch.optim.SGD(lineal_classifier.parameters(),
1e-3,
weight_decay=eval(
self.config['weight_decay']))
epochs = self.config['epochs']
scheduler = torch.optim.lr_scheduler.MultiStepLR(
optimizer,
milestones=[int(0.5 * epochs),
int(0.8 * epochs)],
gamma=0.1,
last_epoch=-1)
criterion = nn.CrossEntropyLoss()
best_acc = 0
for epoch in range(epochs):
print('=' * 20)
print(f'Epoch: {epoch + 1} / {epochs}')
top1 = 0
running_loss = 0
n = 0
lineal_classifier.train()
for idx, (img, lab) in enumerate(train_loader):
B = img.size(0)
img = img.to(self.device)
lab = lab.to(self.device)
optimizer.zero_grad()
loss, top1_batch = self._step_eval_train(
model, lineal_classifier, img, lab)
loss.backward()
optimizer.step()
top1 += (x.argmax(dim=1) == lab).sum().item()
running_loss += loss.item() * B
n += B
print('Training {}/{} - Loss: {:.2f} - top 1: {:.2f}'.format(
idx + 1, len(train_loader), running_loss / n,
100 * top1 / n),
end='\r')
print('\n')
top1 = 0
running_loss = 0
n = 0
lineal_classifier.eval()
for idx, (img, lab) in enumerate(val_loader):
B = img.size(0)
img = img.to(self.device)
lab = lab.to(self.device)
loss, top1_batch = self._step_eval_eval(
model, lineal_classifier, img, lab)
top1 += top1_batch
running_loss += loss.item() * B
n += B
print('Val {}/{} - Loss: {:.2f} - top 1: {:.2f}'.format(
idx + 1, len(val_loader), running_loss / n,
100 * top1 / n),
end='\r')
print('\n')
if best_acc < top1:
best_acc = top1
print(f'Best ACC: {best_acc * 100 / n}')
def train(self):
train_loader, valid_loader = self.dataset.get_data_loaders()
model = ResNetSimCLR(**self.config["model"]) # .to(self.device)
if self.config['n_gpu'] > 1:
device_n = len(eval(self.config['gpu_ids']))
model = torch.nn.DataParallel(model, device_ids=range(device_n))
model = self._load_pre_trained_weights(model)
model = model.to(self.device)
optimizer = torch.optim.Adam(model.parameters(),
1e-5,
weight_decay=eval(
self.config['weight_decay']))
# scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=len(train_loader), eta_min=0,
# last_epoch=-1)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
optimizer, T_max=self.config['epochs'], eta_min=0, last_epoch=-1)
if apex_support and self.config['fp16_precision']:
model, optimizer = amp.initialize(model,
optimizer,
opt_level='O2',
keep_batchnorm_fp32=True)
model_checkpoints_folder = os.path.join(self.writer.log_dir,
'checkpoints')
# save config file
_save_config_file(model_checkpoints_folder)
n_iter = 0
valid_n_iter = 0
best_valid_loss = np.inf
for epoch_counter in range(self.config['epochs']):
for (xis, xjs) in train_loader:
optimizer.zero_grad()
xis = xis.to(self.device)
xjs = xjs.to(self.device)
loss = self._step(model, xis, xjs, n_iter)
if n_iter % self.config['log_every_n_steps'] == 0:
self.writer.add_scalar('train_loss',
loss,
global_step=n_iter)
if apex_support and self.config['fp16_precision']:
with amp.scale_loss(loss, optimizer) as scaled_loss:
scaled_loss.backward()
else:
loss.backward()
optimizer.step()
n_iter += 1
# validate the model if requested
if epoch_counter % self.config['eval_every_n_epochs'] == 0:
valid_loss = self._validate(model, valid_loader)
if valid_loss < best_valid_loss:
# save the model weights
best_valid_loss = valid_loss
torch.save(
model.state_dict(),
os.path.join(model_checkpoints_folder, 'model.pth'))
print('saved')
self.writer.add_scalar('validation_loss',
valid_loss,
global_step=valid_n_iter)
valid_n_iter += 1
# warmup for the first 10 epochs
if epoch_counter >= 10:
scheduler.step()
self.writer.add_scalar('cosine_lr_decay',
scheduler.get_lr()[0],
global_step=n_iter)
def run(config):
model = ResNetSimCLR('resnet50', config.out_dim)
model.load_state_dict(
torch.load(config.model_file, map_location=config.device))
model = model.to(config.device)
clf = nn.Linear(2048, 196)
train_data_dir = os.path.join(config.data_root, 'cars_train/')
train_annos = os.path.join(config.data_root, 'devkit/cars_train_annos.mat')
valid_data_dir = os.path.join(config.data_root, 'cars_test/')
valid_annos = os.path.join(config.data_root,
'devkit/cars_test_annos_withlabels.mat')
if config.encodings_file_prefix:
train_dataset = EncodedStanfordCarsDataset(
train_annos, config.encodings_file_prefix + '-train_encodings.pt')
train_loader = DataLoader(train_dataset,
batch_size=config.encodings_batch_size,
shuffle=True)
valid_dataset = EncodedStanfordCarsDataset(
valid_annos, config.encodings_file_prefix + '-valid_encodings.pt')
valid_loader = DataLoader(valid_dataset,
batch_size=config.encodings_batch_size)
tmp_clf = nn.Linear(2048, 196)
clf_solver = CESolver(tmp_clf,
train_loader,
valid_loader,
config.save_root,
name=config.name + '-clf',
device=config.device)
clf_solver.train(config.encodings_num_epochs)
clf_filename = os.path.join(config.save_root, f'{config.name}-clf.pth')
clf.load_state_dict(
torch.load(clf_filename, map_location=config.device))
full = FineTuner(model, clf)
t = T.Compose([
T.Resize(512),
T.CenterCrop(512),
T.ToTensor(),
T.Lambda(lambda x: x.repeat(3, 1, 1) if x.shape[0] == 1 else x)
])
train_dataset = StanfordCarsMini(train_annos, train_data_dir, t)
train_loader = DataLoader(train_dataset,
batch_size=config.batch_size,
shuffle=True)
valid_dataset = CarsDataset(valid_annos, valid_data_dir, t)
valid_loader = DataLoader(valid_dataset, batch_size=config.batch_size)
solver = CESolver(full,
train_loader,
valid_loader,
config.save_root,
name=config.name,
device=config.device)
# full = full.to(config.device)
# print(solver.validate(full, nn.CrossEntropyLoss()))
solver.train(config.num_epochs)