class BCNNTrainer(object):
"""Manager class to train bilinear CNN.
Attributes:
_options: Hyperparameters.
_path: Useful paths.
_net: Bilinear CNN.
_criterion: Cross-entropy loss.
_solver: SGD with momentum.
_scheduler: Reduce learning rate by a fator of 0.1 when plateau.
_train_loader: Training data.
_test_loader: Testing data.
"""
def __init__(self, options, path, ckpt_basename='vgg_16'):
"""Prepare the network, criterion, solver, and data.
Args:
options, dict: Hyperparameters.
"""
print('Prepare the network and data.')
self._options = options
self._path = path
self.ckpt_basename = ckpt_basename
# Network.
self._net = BCNN(freeze_features=True)
#self._net = torch.nn.DataParallel(self._net)
self._net.features = torch.nn.DataParallel(self._net.features)
self._net.cuda()
if 'ckpt_path' in self._path:
if os.path.exists(self._path['ckpt_path']):
print('Continue from', self._path['ckpt_path'])
self._net.load_state_dict(torch.load(self._path['ckpt_path']))
else:
print('Ckpt {} not found!'.format(self._path['ckpt_path']))
print(self._net)
# Criterion.
self._criterion = torch.nn.CrossEntropyLoss().cuda()
# Solver.
self._solver = torch.optim.SGD(
self._net.fc.parameters(),
lr=self._options['base_lr'],
momentum=0.9,
weight_decay=self._options['weight_decay'])
if self._options['lr_scheduler'] == 'reduce_on_plateau':
self._scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
self._solver,
mode='max',
factor=0.1,
patience=5,
verbose=True,
threshold=1e-4,
min_lr=1e-6)
elif self._options['lr_scheduler'] == 'fixed':
self._scheduler = torch.optim.lr_scheduler.LambdaLR(
self._solver, lambda epoch: 1.0)
else:
raise ValueError('Unknown scheduler:',
self._options['lr_scheduler'])
# Imagenet normalization
normalize = torchvision.transforms.Normalize(
mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
train_transforms = torchvision.transforms.Compose([
torchvision.transforms.Resize(size=448), # Let smaller edge match
torchvision.transforms.RandomHorizontalFlip(),
torchvision.transforms.RandomCrop(size=448),
torchvision.transforms.ToTensor(),
normalize
])
test_transforms = torchvision.transforms.Compose([
torchvision.transforms.Resize(size=448),
torchvision.transforms.CenterCrop(size=448),
torchvision.transforms.ToTensor(), normalize
])
train_data = cub200.CUB200(root=self._path['cub200'],
train=True,
download=True,
transform=train_transforms)
test_data = cub200.CUB200(root=self._path['cub200'],
train=False,
download=True,
transform=test_transforms)
self._train_loader = torch.utils.data.DataLoader(
train_data,
batch_size=self._options['batch_size'],
shuffle=True,
num_workers=4,
pin_memory=True)
self._test_loader = torch.utils.data.DataLoader(test_data,
batch_size=16,
shuffle=False,
num_workers=4,
pin_memory=True)
def train(self):
"""Train the network."""
print('Training.')
self._net.train()
best_acc = 0.0
best_epoch = None
for epoch in range(self._options['epochs']):
epoch_loss = []
num_correct = 0
num_total = 0
for batch_idx, (X, y) in enumerate(self._train_loader):
# Data.
X = X.cuda(non_blocking=True)
y = y.cuda(non_blocking=True)
# Clear the existing gradients.
self._solver.zero_grad()
# Forward pass.
score = self._net(X)
loss = self._criterion(score, y)
epoch_loss.append(loss.data.item())
# Prediction.
_, prediction = torch.max(score.data, 1)
num_total += y.size(0)
num_correct += torch.sum(prediction == y.data).item()
# Backward pass.
loss.backward()
self._solver.step()
sys.stdout.write('\r')
sys.stdout.write(
'| Epoch [%3d/%3d] Iter[%3d/%3d]\t\tLoss: %.4f Acc@1: %.3f%%'
% (epoch, self._options['epochs'], batch_idx + 1,
len(self._train_loader), loss.data.item(),
(100. * num_correct) / num_total))
sys.stdout.flush()
train_acc = 100 * num_correct / num_total
test_acc = self._accuracy(self._test_loader)
print('\nEpoch\tTrain loss\tTrain acc\tTest acc')
print('%d\t%4.3f\t\t%4.2f%%\t\t%4.2f%%' %
(epoch + 1, np.mean(epoch_loss), train_acc, test_acc))
self._scheduler.step(test_acc)
if test_acc > best_acc:
best_acc = test_acc
best_epoch = epoch + 1
print('*', end='')
# Save model onto disk.
save_path = os.path.join(
self._path['model'],
'{}_epoch_{}.pth'.format(self.ckpt_basename, epoch + 1))
save_path_best = os.path.join(
self._path['model'],
'{}_epoch_best.pth'.format(self.ckpt_basename))
torch.save(self._net.state_dict(), save_path)
shutil.copy(save_path, save_path_best)
print('Best at epoch %d, test accuaray %f' % (best_epoch, best_acc))
def _accuracy(self, data_loader):
"""Compute the train/test accuracy.
Args:
data_loader: Train/Test DataLoader.
Returns:
Train/Test accuracy in percentage.
"""
self._net.eval()
num_correct = 0
num_total = 0
for X, y in data_loader:
# Data.
X = X.cuda(non_blocking=True)
y = y.cuda(non_blocking=True)
with torch.no_grad():
# Prediction.
score = self._net(X)
_, prediction = torch.max(score.data, 1)
num_total += y.size(0)
num_correct += torch.sum(prediction == y.data).item()
return 100 * num_correct / num_total
def getStat(self):
"""Get the mean and std value for a certain dataset."""
print('Compute mean and variance for training data.')
train_data = cub200.CUB200(root=self._path['cub200'],
train=True,
transform=torchvision.transforms.ToTensor(),
download=True)
train_loader = torch.utils.data.DataLoader(train_data,
batch_size=1,
shuffle=False,
num_workers=4,
pin_memory=True)
mean = torch.zeros(3)
std = torch.zeros(3)
for X, _ in tqdm(train_loader):
for d in range(3):
mean[d] += X[:, d, :, :].mean()
std[d] += X[:, d, :, :].std()
mean.div_(len(train_data))
std.div_(len(train_data))
print(mean)
print(std)
class BCNNManager(object):
"""Manager class to train bilinear CNN.
Attributes:
_options: Hyperparameters.
_path: Useful paths.
_net: Bilinear CNN.
_criterion: Cross-entropy loss.
_solver: SGD with momentum.
_scheduler: Reduce learning rate by a fator of 0.1 when plateau.
_train_loader: Training data.
_test_loader: Testing data.
"""
def __init__(self, options, path):
"""Prepare the network, criterion, solver, and data.
Args:
options, dict: Hyperparameters.
"""
print('Prepare the network and data.')
self._options = options
self._path = path
# Network.
if self._options['dataset'] == 'cub200':
num_classes = 200
elif self._options['dataset'] == 'aircraft':
num_classes = 100
else:
raise NotImplementedError("Dataset " + self._options['dataset'] +
" is not implemented.")
self._net = BCNN(num_classes=num_classes,
pretrained=options['target'] == 'fc')
# Load the model from disk.
if options['target'] == 'all':
self._net.load_state_dict(torch.load(self._path['model']))
self._net = torch.nn.parallel.DistributedDataParallel(
self._net.cuda(),
device_ids=[self._options['local_rank']],
output_device=self._options['local_rank'])
if dist.get_rank() == 0:
print(self._net)
# Criterion.
self._criterion = torch.nn.CrossEntropyLoss().cuda()
# Solver.
self._solver = torch.optim.SGD(
self._net.module.trainable_params,
lr=self._options['base_lr'] * dist.get_world_size(),
momentum=0.9,
weight_decay=self._options['weight_decay'])
self._scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
self._solver,
mode='max',
factor=0.1,
patience=3,
verbose=True,
threshold=1e-4)
train_transforms = torchvision.transforms.Compose([
torchvision.transforms.Resize(size=448), # Let smaller edge match
torchvision.transforms.RandomHorizontalFlip(),
torchvision.transforms.RandomCrop(size=448),
torchvision.transforms.ToTensor(),
torchvision.transforms.Normalize(mean=(0.485, 0.456, 0.406),
std=(0.229, 0.224, 0.225))
])
test_transforms = torchvision.transforms.Compose([
torchvision.transforms.Resize(size=448),
torchvision.transforms.CenterCrop(size=448),
torchvision.transforms.ToTensor(),
torchvision.transforms.Normalize(mean=(0.485, 0.456, 0.406),
std=(0.229, 0.224, 0.225))
])
if self._options['dataset'] == 'cub200':
train_data = cub200.CUB200(root=self._path['dataset'],
train=True,
download=True,
transform=train_transforms)
test_data = cub200.CUB200(root=self._path['dataset'],
train=False,
download=True,
transform=test_transforms)
elif self._options['dataset'] == 'aircraft':
train_data = aircraft.Aircraft(root=self._path['dataset'],
train=True,
download=True,
transform=train_transforms)
test_data = aircraft.Aircraft(root=self._path['dataset'],
train=False,
download=True,
transform=test_transforms)
else:
raise NotImplementedError("Dataset " + self._options['dataset'] +
" is not implemented.")
# Partition dataset among workers using DistributedSampler
train_sampler = distributed.DistributedSampler(
train_data,
num_replicas=dist.get_world_size(),
rank=dist.get_rank())
test_sampler = distributed.DistributedSampler(
test_data,
num_replicas=dist.get_world_size(),
rank=dist.get_rank())
self._train_loader = torch.utils.data.DataLoader(
train_data,
batch_size=self._options['batch_size'],
shuffle=False,
num_workers=4,
pin_memory=True,
sampler=train_sampler)
self._test_loader = torch.utils.data.DataLoader(
test_data,
batch_size=self._options['batch_size'],
shuffle=False,
num_workers=4,
pin_memory=True,
sampler=test_sampler)
def train(self):
"""Train the network."""
best_acc = 0.0
best_epoch = None
if dist.get_rank() == 0:
print('Training.')
print('Epoch\tTrain loss\tTrain acc\tTest acc\tTrain time')
for t in range(self._options['epochs']):
t0 = time.time()
self._train_loader.sampler.set_epoch(t)
epoch_loss = []
num_correct = 0
num_total = 0
for X, y in self._train_loader:
# Data.
X = torch.autograd.Variable(X.cuda())
y = torch.autograd.Variable(y.cuda(async=True))
# Clear the existing gradients.
self._solver.zero_grad()
# Forward pass.
score = self._net(X)
loss = self._criterion(score, y)
epoch_loss.append(loss.item())
# Prediction.
_, prediction = torch.max(score.data, 1)
num_total += y.size(0)
num_correct += torch.sum(prediction == y.data)
# Backward pass.
loss.backward()
self._solver.step()
test_acc = self._accuracy(self._test_loader)
self._scheduler.step(test_acc)
if dist.get_rank() == 0:
train_acc = 100 * num_correct / num_total
if test_acc > best_acc:
best_acc = test_acc
best_epoch = t + 1
print('*', end='')
# Save model onto disk.
torch.save(
self._net.module.state_dict(),
os.path.join(self._path['model_dir'],
'vgg_16_epoch_%d.pth' % (t + 1)))
print('%d\t%4.3f\t\t%4.2f%%\t\t%4.2f%%\t\t%4.2fs' %
(t + 1, sum(epoch_loss) / len(epoch_loss), train_acc,
test_acc, time.time() - t0))
if dist.get_rank() == 0:
print('Best at epoch %d, test accuaray %f' %
(best_epoch, best_acc))
def _accuracy(self, data_loader):
"""Compute the train/test accuracy.
Args:
data_loader: Train/Test DataLoader.
Returns:
Train/Test accuracy in percentage.
"""
self._net.train(False)
num_correct = 0
num_total = 0
with torch.no_grad():
for X, y in data_loader:
# Data.
X = torch.autograd.Variable(X.cuda())
y = torch.autograd.Variable(y.cuda(async=True))
# Prediction.
score = self._net(X)
_, prediction = torch.max(score.data, 1)
num_total += y.size(0)
num_correct += torch.sum(prediction == y.data).item()
self._net.train(True) # Set the model to training phase
num_total = torch.tensor(num_total).cuda()
num_correct = torch.tensor(num_correct).cuda()
dist.all_reduce(num_total, op=dist.ReduceOp.SUM)
dist.all_reduce(num_correct, op=dist.ReduceOp.SUM)
return 100 * num_correct.data.item() / num_total.data.item()
def getStat(self):
"""Get the mean and std value for a certain dataset."""
print('Compute mean and variance for training data.')
train_data = cub200.CUB200(root=self._path['cub200'],
train=True,
transform=torchvision.transforms.ToTensor(),
download=True)
train_loader = torch.utils.data.DataLoader(train_data,
batch_size=1,
shuffle=False,
num_workers=4,
pin_memory=True)
mean = torch.zeros(3)
std = torch.zeros(3)
for X, _ in train_loader:
for d in range(3):
mean[d] += X[:, d, :, :].mean()
std[d] += X[:, d, :, :].std()
mean.div_(len(train_data))
std.div_(len(train_data))
print(mean)
print(std)