Spaces:
Runtime error
Runtime error
File size: 6,090 Bytes
54660f7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 |
import torch
from models.resnet_simclr import ResNetSimCLR
from torch.utils.tensorboard import SummaryWriter
import torch.nn.functional as F
from loss.nt_xent import NTXentLoss
import os
import shutil
import sys
apex_support = False
try:
sys.path.append('./apex')
from apex import amp
apex_support = True
except:
print("Please install apex for mixed precision training from: https://github.com/NVIDIA/apex")
apex_support = False
import numpy as np
torch.manual_seed(0)
def _save_config_file(model_checkpoints_folder):
if not os.path.exists(model_checkpoints_folder):
os.makedirs(model_checkpoints_folder)
shutil.copy('./config.yaml', os.path.join(model_checkpoints_folder, 'config.yaml'))
class SimCLR(object):
def __init__(self, dataset, config, args=None):
self.config = config
self.device = self._get_device()
self.writer = SummaryWriter()
self.dataset = dataset
self.nt_xent_criterion = NTXentLoss(self.device, config['batch_size'], **config['loss'])
self.args = args
def _get_device(self):
device = 'cuda' if torch.cuda.is_available() else 'cpu'
print("Running on:", device)
return device
def _step(self, model, xis, xjs, n_iter):
# get the representations and the projections
ris, zis = model(xis) # [N,C]
# get the representations and the projections
rjs, zjs = model(xjs) # [N,C]
# normalize projection feature vectors
zis = F.normalize(zis, dim=1)
zjs = F.normalize(zjs, dim=1)
loss = self.nt_xent_criterion(zis, zjs)
return loss
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:
model = torch.nn.DataParallel(model, device_ids=eval(self.config['gpu_ids']))
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)
if self.args is None:
model_checkpoints_folder = os.path.join(self.writer.log_dir, 'checkpoints')
else:
model_checkpoints_folder = self.args.dest_weights#os.environ['FEATURE_EXTRACTOR_WEIGHT_PATH']
model_checkpoints_folder = os.path.dirname(model_checkpoints_folder)
# 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)
print("[%d/%d] step: %d train_loss: %.3f" % (epoch_counter, self.config['epochs'], n_iter, loss))
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)
print("[%d/%d] val_loss: %.3f" % (epoch_counter, self.config['epochs'], valid_loss))
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 _load_pre_trained_weights(self, model):
try:
checkpoints_folder = os.path.join('./runs', self.config['fine_tune_from'], 'checkpoints')
state_dict = torch.load(os.path.join(checkpoints_folder, 'model.pth'))
model.load_state_dict(state_dict)
print("Loaded pre-trained model with success.")
except FileNotFoundError:
print("Pre-trained weights not found. Training from scratch.")
return model
def _validate(self, model, valid_loader):
# validation steps
with torch.no_grad():
model.eval()
valid_loss = 0.0
counter = 0
for (xis, xjs) in valid_loader:
xis = xis.to(self.device)
xjs = xjs.to(self.device)
loss = self._step(model, xis, xjs, counter)
valid_loss += loss.item()
counter += 1
valid_loss /= counter
model.train()
return valid_loss
|