Hi guys, I saw all the posts about out of memory and tried a lot of things but nothing works… I am using pytorch 1 and this is my code:
#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""Fine-tune the fc layer only for bilinear CNN.
Usage:
CUDA_VISIBLE_DEVICES=0,1,2,3 ./src/bilinear_cnn_fc.py --base_lr 0.05 \
--batch_size 64 --epochs 100 --weight_decay 5e-4
"""
import os
import torch
import torchvision
from torch.utils.data import DataLoader, Dataset
import data_generator as whales_generator
import albumentations
torch.manual_seed(0)
torch.cuda.manual_seed_all(0)
torch.set_default_dtype(torch.float32)
torch.set_default_tensor_type(torch.FloatTensor)
torch.manual_seed(0)
torch.cuda.manual_seed_all(0)
torch.backends.cudnn.benchmark = True
__all__ = ['BCNN', 'BCNNManager']
__author__ = 'Hao Zhang'
__copyright__ = '2018 LAMDA'
__date__ = '2018-01-09'
__email__ = 'zhangh0214@gmail.com'
__license__ = 'CC BY-SA 3.0'
__status__ = 'Development'
__updated__ = '2018-01-13'
__version__ = '1.2'
class BCNN(torch.nn.Module):
"""B-CNN for WHALES.
The B-CNN model is illustrated as follows.
conv1^2 (64) -> pool1 -> conv2^2 (128) -> pool2 -> conv3^3 (256) -> pool3
-> conv4^3 (512) -> pool4 -> conv5^3 (512) -> bilinear pooling
-> sqrt-normalize -> L2-normalize -> fc (200).
The network accepts a 3*448*448 input, and the pool5 activation has shape
512*28*28 since we down-sample 5 times.
Attributes:
features, torch.nn.Module: Convolution and pooling layers.
fc, torch.nn.Module: 200.
"""
def __init__(self):
"""Declare all needed layers."""
torch.nn.Module.__init__(self)
# Convolution and pooling layers of VGG-16.
self.features = torchvision.models.vgg16(pretrained=True).features
self.features = torch.nn.Sequential(*list(self.features.children())[0:16]) # Hay jupyter notebook para mirar
# Linear classifier.
self.fc = torch.nn.Linear(256**2, 5004)
# Freeze all previous layers.
for param in self.features.parameters():
param.requires_grad = False
# Initialize the fc layers.
torch.nn.init.kaiming_normal_(self.fc.weight.data)
if self.fc.bias is not None:
torch.nn.init.constant_(self.fc.bias.data, val=0)
def forward(self, X):
"""Forward pass of the network.
Args:
X, torch.autograd.Variable of shape N*3*448*448.
Returns:
Score, torch.autograd.Variable of shape N*200.
"""
N = X.size()[0]
assert X.size() == (N, 3, 224, 224)
X = self.features(X)
assert X.size() == (N, 256, 56, 56)
X = X.view(N, 256, 56**2)
X = torch.bmm(X, torch.transpose(X, 1, 2)) / (56**2) # Bilinear
assert X.size() == (N, 256, 256)
X = X.view(N, 256**2)
X = torch.sqrt(X + 1e-5)
X = torch.nn.functional.normalize(X)
X = self.fc(X)
#assert X.size() == (N, 200)
return X
class BCNNManager(object):
"""Manager class to train bilinear CNN.
Attributes:
_options: Hyperparameters.
_path: Useful paths.
_net: Bilinear CNN.
_criterion: Cross-entropy loss.
_optimizer: 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.
self._net = torch.nn.DataParallel(BCNN()).cuda()
print(self._net)
# Criterion.
self._criterion = torch.nn.CrossEntropyLoss().cuda()
# Solver.
self._optimizer = torch.optim.SGD(self._net.module.fc.parameters(), lr=self._options['base_lr'],
momentum=0.9, weight_decay=self._options['weight_decay'])
self._scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
self._optimizer, mode='max', factor=0.1, patience=3, verbose=True,
threshold=1e-4)
data_transforms_train = albumentations.Compose([
albumentations.Resize(224 ,224),
albumentations.HorizontalFlip(),
albumentations.Blur(blur_limit=2),
albumentations.ShiftScaleRotate(rotate_limit=7, scale_limit=0.085),
albumentations.RandomBrightness(limit=0.12),
albumentations.Normalize(mean=(0.485, 0.456, 0.406),
std=(0.229, 0.224, 0.225))
])
train_dataset = whales_generator.WhaleDataset(
datatype='train',
transform=data_transforms_train,
normalization="", # No porque ya lo normalizamos con la transformacion
color=True, # Necasario para modelo preentrenado
use_one=False, # Queremos todas las ballenitas!
humpback_only=True # Solo las colitas
)
self._train_loader = DataLoader(train_dataset, batch_size=self._options['batch_size'], pin_memory=False, shuffle=True)
"""
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))
])
test_data = cub200.CUB200(
root=self._path['cub200'], train=False, download=True,
transform=test_transforms)
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.')
best_acc = 0.0
best_epoch = None
print('Epoch\tTrain loss\tTrain acc\tTest acc')
for t in range(self._options['epochs']):
epoch_loss = 0
num_correct = 0
num_total = 0
total_batches = 0
for X, y in self._train_loader:
# Data.
# Clear the existing gradients.
X = X.cuda()
y = y.cuda()
# Forward pass.
score = self._net(X)
loss = self._criterion(score, y.long())
with torch.no_grad():
epoch_loss += loss.item()
# Prediction.
prediction = torch.argmax(score, dim=1)
num_total += y.size(0)
num_correct += torch.sum(prediction == y.long()).item()
# Backward pass.
self._optimizer.zero_grad()
loss.backward()
self._optimizer.step()
total_batches+=1
del X, y, score, loss, prediction
train_acc = 100 * num_correct / num_total
#test_acc = self._accuracy(self._test_loader)
self._scheduler.step(test_acc)
if test_acc > best_acc:
best_acc = test_acc
best_epoch = t + 1
print('*', end='')
# Save model onto disk.
torch.save(self._net.state_dict(),
os.path.join(self._path['model'],
'vgg_16_epoch_%d.pth' % (t + 1)))
print('%d\t%4.3f\t\t%4.2f%%\t\t%4.2f%%' %
(t+1, epoch_loss / total_batches, train_acc, test_acc))
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
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.logn().data).item()
self._net.train(True) # Set the model to training phase
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.')
data_transforms_train = albumentations.Compose([
albumentations.Resize(448, 448)
])
train_data = whales_generator.WhaleDataset(
datatype='train',
transform=data_transforms_train,
normalization=norm_data,
color=color,
use_one=use_one,
humpback_only=True
)
train_loader = DataLoader(train_data, batch_size=1, pin_memory=True, shuffle=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)
def main():
"""The main function."""
import argparse
parser = argparse.ArgumentParser(
description='Train bilinear CNN on WHALES.')
"""
parser.add_argument('--base_lr', dest='base_lr', type=float, required=True,
help='Base learning rate for training.')
parser.add_argument('--batch_size', dest='batch_size', type=int,
required=True, help='Batch size.')
parser.add_argument('--epochs', dest='epochs', type=int,
required=True, help='Epochs for training.')
parser.add_argument('--weight_decay', dest='weight_decay', type=float,
required=True, help='Weight decay.')
"""
parser.add_argument('--base_lr', dest='base_lr', type=float, required=False, default=0.1,
help='Base learning rate for training.')
parser.add_argument('--batch_size', dest='batch_size', type=int, default=2,
required=False, help='Batch size.')
parser.add_argument('--epochs', dest='epochs', type=int, default=10,
required=False, help='Epochs for training.')
parser.add_argument('--weight_decay', dest='weight_decay', type=float, default=1e-8,
required=False, help='Weight decay.')
args = parser.parse_args()
if args.base_lr <= 0:
raise AttributeError('--base_lr parameter must >0.')
if args.batch_size <= 0:
raise AttributeError('--batch_size parameter must >0.')
if args.epochs < 0:
raise AttributeError('--epochs parameter must >=0.')
if args.weight_decay <= 0:
raise AttributeError('--weight_decay parameter must >0.')
options = {
'base_lr': args.base_lr,
'batch_size': args.batch_size,
'epochs': args.epochs,
'weight_decay': args.weight_decay,
}
project_root = os.popen('pwd').read().strip()
path = {
'model': os.path.join(project_root + "/mario/bilinear", 'checkpoints'),
}
for d in path:
assert os.path.isdir(path[d])
manager = BCNNManager(options, path)
# manager.getStat()
manager.train()
if __name__ == '__main__':
main()
When starts, epoch 1 consumes 5Gb but on second epoch it runs out of memory at backward point 