The following script is basically 99% from NVIDIAs Jetson AI github, which uses Pytorch to train different DNNs for Image Classification. Usually the scripts expect a command line argument to provide the path to training data. I changed this to use the Fashion MNIST dataset (line 156f).
The problem is, that the script expects RGB images, so the output shape is [3,x,y], but Fashion MNIST is grayscale and has only one channel. Is there a way to change this in the script or do I need to modify the used pretrained model?
import argparse
import os
import random
import time
import shutil
import warnings
import datetime
import torch
import torch.nn as nn
import torch.nn.parallel
import torch.nn.functional as F
import torch.backends.cudnn as cudnn
import torch.optim
import torch.utils.data
import torchvision.transforms as transforms
import torchvision.datasets as datasets
import torchvision.models as models
from torch.utils.tensorboard import SummaryWriter
from voc import VOCDataset
from nuswide import NUSWideDataset
from reshape import reshape_model
# get the available network architectures
model_names = sorted(name for name in models.__dict__
if name.islower() and not name.startswith("__")
and callable(models.__dict__[name]))
# parse command-line arguments
parser = argparse.ArgumentParser(description='PyTorch Image Classifier Training')
# Das wahrscheinlich raus und durch Laden von Fashion MNIST ersetzen
#parser.add_argument('data', metavar='DIR',
# help='path to dataset')
parser.add_argument('--dataset-type', type=str, default='folder',
choices=['folder', 'nuswide', 'voc'],
help='specify the dataset type (default: folder)')
parser.add_argument('--multi-label', action='store_true',
help='multi-label model (aka image tagging)')
parser.add_argument('--multi-label-threshold', type=float, default=0.5,
help='confidence threshold for counting a prediction as correct')
parser.add_argument('--model-dir', type=str, default='models',
help='path to desired output directory for saving model '
'checkpoints (default: models/)')
parser.add_argument('-a', '--arch', metavar='ARCH', default='resnet18',
choices=model_names,
help='model architecture: ' + ' | '.join(model_names) + ' (default: resnet18)')
parser.add_argument('--resolution', default=224, type=int, metavar='N',
help='input NxN image resolution of model (default: 224x224) '
'note than Inception models should use 299x299')
parser.add_argument('-j', '--workers', default=2, type=int, metavar='N',
help='number of data loading workers (default: 2)')
parser.add_argument('--epochs', default=35, type=int, metavar='N',
help='number of total epochs to run')
parser.add_argument('--start-epoch', default=0, type=int, metavar='N',
help='manual epoch number (useful on restarts)')
parser.add_argument('-b', '--batch-size', default=8, type=int, metavar='N',
help='mini-batch size (default: 8)')
parser.add_argument('--lr', '--learning-rate', default=0.1, type=float,
metavar='LR', help='initial learning rate', dest='lr')
parser.add_argument('--momentum', default=0.9, type=float, metavar='M',
help='momentum')
parser.add_argument('--wd', '--weight-decay', default=1e-4, type=float,
metavar='W', help='weight decay (default: 1e-4)',
dest='weight_decay')
parser.add_argument('-p', '--print-freq', default=10, type=int,
metavar='N', help='print frequency (default: 10)')
parser.add_argument('--resume', default='', type=str, metavar='PATH',
help='path to latest checkpoint (default: none)')
parser.add_argument('-e', '--evaluate', dest='evaluate', action='store_true',
help='evaluate model on validation set')
parser.add_argument('--pretrained', dest='pretrained', action='store_true', default=True,
help='use pre-trained model')
parser.add_argument('--seed', default=None, type=int,
help='seed for initializing training')
parser.add_argument('--gpu', default=0, type=int,
help='GPU ID to use (default: 0)')
args = parser.parse_args()
# open tensorboard logger (to model_dir/tensorboard)
tensorboard = SummaryWriter(log_dir=os.path.join(args.model_dir, "tensorboard", f"{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}"))
print(f"To start tensorboard run: tensorboard --log-dir={os.path.join(args.model_dir, 'tensorboard')}")
# variable for storing the best model accuracy so far
best_accuracy = 0
def main(args):
"""
Load dataset, setup model, and train for N epochs
"""
global best_accuracy
if args.seed is not None:
random.seed(args.seed)
torch.manual_seed(args.seed)
cudnn.deterministic = True
warnings.warn('You have chosen to seed training. '
'This will turn on the CUDNN deterministic setting, '
'which can slow down your training considerably! '
'You may see unexpected behavior when restarting '
'from checkpoints.')
if args.gpu is not None:
print(f"=> using GPU {args.gpu} ({torch.cuda.get_device_name(args.gpu)})")
# setup data transformations
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
train_transforms = transforms.Compose([
transforms.RandomResizedCrop(args.resolution),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize,
])
val_transforms = transforms.Compose([
transforms.Resize(args.resolution),
transforms.CenterCrop(args.resolution),
transforms.ToTensor(),
normalize,
])
# load the dataset
#
##if args.dataset_type == 'folder':
## train_dataset = datasets.ImageFolder(os.path.join(args.data, 'train'), train_transforms)
## val_dataset = datasets.ImageFolder(os.path.join(args.data, 'val'), val_transforms)
##elif args.dataset_type == 'nuswide':
## train_dataset = NUSWideDataset(args.data, 'trainval', train_transforms)
## val_dataset = NUSWideDataset(args.data, 'test', val_transforms)
##elif args.dataset_type == 'voc':
## train_dataset = VOCDataset(args.data, 'trainval', train_transforms)
## val_dataset = VOCDataset(args.data, 'val', val_transforms)
train_dataset = datasets.FashionMNIST('~/.pytorch/F_MNIST_data', download=True, train=True, transform=train_transforms)
val_dataset = datasets.FashionMNIST('~/.pytorch/F_MNIST_data', download=True, train=False, transform=val_transforms)
if (args.dataset_type == 'nuswide' or args.dataset_type == 'voc') and (not args.multi_label):
raise ValueError("nuswide or voc datasets should be run with --multi-label")
print(f"=> dataset classes: {len(train_dataset.classes)} {train_dataset.classes}")
train_loader = torch.utils.data.DataLoader(
train_dataset, batch_size=args.batch_size, shuffle=True,
num_workers=args.workers, pin_memory=True)
val_loader = torch.utils.data.DataLoader(
val_dataset, batch_size=args.batch_size, shuffle=False,
num_workers=args.workers, pin_memory=True)
# create or load the model if using pre-trained (the default)
if args.pretrained:
print(f"=> using pre-trained model '{args.arch}'")
model = models.__dict__[args.arch](pretrained=True)
else:
print(f"=> creating model '{args.arch}'")
model = models.__dict__[args.arch]()
# reshape the model for the number of classes in the dataset
model = reshape_model(model, args.arch, len(train_dataset.classes))
# define loss function (criterion) and optimizer
if args.multi_label:
criterion = nn.BCEWithLogitsLoss()
else:
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
# transfer the model to the GPU that it should be run on
if args.gpu is not None:
torch.cuda.set_device(args.gpu)
model = model.cuda(args.gpu)
criterion = criterion.cuda(args.gpu)
# optionally resume from a checkpoint
if args.resume:
if os.path.isfile(args.resume):
print(f"=> loading checkpoint '{args.resume}'")
checkpoint = torch.load(args.resume)
args.start_epoch = checkpoint['epoch'] + 1
#best_accuracy = checkpoint['best_accuracy']
#if args.gpu is not None:
# best_accuracy = best_accuracy.to(args.gpu) # best_accuracy may be from a checkpoint from a different GPU
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
print(f"=> loaded checkpoint '{args.resume}' (epoch {checkpoint['epoch']})")
else:
print(f"=> no checkpoint found at '{args.resume}'")
cudnn.benchmark = True
# if in evaluation mode, only run validation
if args.evaluate:
validate(val_loader, model, criterion, 0)
return
# train for the specified number of epochs
for epoch in range(args.start_epoch, args.epochs):
# decay the learning rate
adjust_learning_rate(optimizer, epoch)
# train for one epoch
train_loss, train_acc = train(train_loader, model, criterion, optimizer, epoch)
# evaluate on validation set
val_loss, val_acc = validate(val_loader, model, criterion, epoch)
# remember best acc@1 and save checkpoint
is_best = val_acc > best_accuracy
best_accuracy = max(val_acc, best_accuracy)
print(f"=> Epoch {epoch}")
print(f" * Train Loss {train_loss:.4e}")
print(f" * Train Accuracy {train_acc:.4f}")
print(f" * Val Loss {val_loss:.4e}")
print(f" * Val Accuracy {val_acc:.4f}{'*' if is_best else ''}")
save_checkpoint({
'epoch': epoch,
'arch': args.arch,
'resolution': args.resolution,
'classes': train_dataset.classes,
'num_classes': len(train_dataset.classes),
'multi_label': args.multi_label,
'state_dict': model.state_dict(),
'accuracy': {'train': train_acc, 'val': val_acc},
'loss' : {'train': train_loss, 'val': val_loss},
'optimizer' : optimizer.state_dict(),
}, is_best)
def train(train_loader, model, criterion, optimizer, epoch):
"""
Train one epoch over the dataset
"""
batch_time = AverageMeter('Time', ':6.3f')
data_time = AverageMeter('Data', ':6.3f')
losses = AverageMeter('Loss', ':.4e')
acc = AverageMeter('Accuracy', ':7.3f')
progress = ProgressMeter(
len(train_loader),
[batch_time, data_time, losses, acc],
prefix=f"Epoch: [{epoch}]")
# switch to train mode
model.train()
# get the start time
epoch_start = time.time()
end = epoch_start
# train over each image batch from the dataset
for i, (images, target) in enumerate(train_loader):
# measure data loading time
data_time.update(time.time() - end)
if args.gpu is not None:
images = images.cuda(args.gpu, non_blocking=True)
target = target.cuda(args.gpu, non_blocking=True)
# compute output
output = model(images)
loss = criterion(output, target)
# record loss and measure accuracy
losses.update(loss.item(), images.size(0))
acc.update(accuracy(output, target), images.size(0))
# compute gradient and do SGD step
optimizer.zero_grad()
loss.backward()
optimizer.step()
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
if i % args.print_freq == 0 or i == len(train_loader)-1:
progress.display(i)
print(f"Epoch: [{epoch}] completed, elapsed time {time.time() - epoch_start:6.3f} seconds")
tensorboard.add_scalar('Loss/train', losses.avg, epoch)
tensorboard.add_scalar('Accuracy/train', acc.avg, epoch)
return losses.avg, acc.avg
def validate(val_loader, model, criterion, epoch):
"""
Measure model performance across the val dataset
"""
batch_time = AverageMeter('Time', ':6.3f')
losses = AverageMeter('Loss', ':.4e')
acc = AverageMeter('Accuracy', ':7.3f')
progress = ProgressMeter(
len(val_loader),
[batch_time, losses, acc],
prefix='Val: ')
# switch to evaluate mode
model.eval()
with torch.no_grad():
end = time.time()
for i, (images, target) in enumerate(val_loader):
if args.gpu is not None:
images = images.cuda(args.gpu, non_blocking=True)
target = target.cuda(args.gpu, non_blocking=True)
# compute output
output = model(images)
loss = criterion(output, target)
# record loss and measure accuracy
losses.update(loss.item(), images.size(0))
acc.update(accuracy(output, target), images.size(0))
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
if i % args.print_freq == 0 or i == len(val_loader)-1:
progress.display(i)
tensorboard.add_scalar('Loss/val', losses.avg, epoch)
tensorboard.add_scalar('Accuracy/val', acc.avg, epoch)
return losses.avg, acc.avg
def save_checkpoint(state, is_best, filename='checkpoint.pth.tar', best_filename='model_best.pth.tar', labels_filename='labels.txt'):
"""
Save a model checkpoint file, along with the best-performing model if applicable
"""
if args.model_dir:
model_dir = os.path.expanduser(args.model_dir)
if not os.path.exists(model_dir):
os.mkdir(model_dir)
filename = os.path.join(model_dir, filename)
best_filename = os.path.join(model_dir, best_filename)
labels_filename = os.path.join(model_dir, labels_filename)
# save the checkpoint
torch.save(state, filename)
# earmark the best checkpoint
if is_best:
shutil.copyfile(filename, best_filename)
print(f"saved best model to: {best_filename}")
else:
print(f"saved checkpoint to: {filename}")
# save labels.txt on the first epoch
if state['epoch'] == 0:
with open(labels_filename, 'w') as file:
for label in state['classes']:
file.write(f"{label}\n")
print(f"saved class labels to: {labels_filename}")
def adjust_learning_rate(optimizer, epoch):
"""
Sets the learning rate to the initial LR decayed by 10 every 30 epochs
"""
lr = args.lr * (0.1 ** (epoch // 30))
for param_group in optimizer.param_groups:
param_group['lr'] = lr
def accuracy(output, target):
"""
Computes the accuracy of predictions vs groundtruth
"""
with torch.no_grad():
if args.multi_label:
output = F.sigmoid(output)
preds = ((output >= args.multi_label_threshold) == target.bool()) # https://medium.com/@yrodriguezmd/tackling-the-accuracy-multi-metric-9e2356f62513
# https://stackoverflow.com/a/61585551
#output[output >= args.multi_label_threshold] = 1
#output[output < args.multi_label_threshold] = 0
#preds = (output == target)
else:
output = F.softmax(output, dim=-1)
_, preds = torch.max(output, dim=-1)
preds = (preds == target)
return preds.float().mean().cpu().item() * 100.0
class AverageMeter(object):
"""
Computes and stores the average and current value
"""
def __init__(self, name, fmt=':f'):
self.name = name
self.fmt = fmt
self.reset()
def reset(self):
self.val = 0
self.avg = 0
self.sum = 0
self.count = 0
def update(self, val, n=1):
self.val = val
self.sum += val * n
self.count += n
self.avg = self.sum / self.count
def __str__(self):
fmtstr = '{name} {val' + self.fmt + '} ({avg' + self.fmt + '})'
return fmtstr.format(**self.__dict__)
class ProgressMeter(object):
"""
Progress metering
"""
def __init__(self, num_batches, meters, prefix=""):
self.batch_fmtstr = self._get_batch_fmtstr(num_batches)
self.meters = meters
self.prefix = prefix
def display(self, batch):
entries = [self.prefix + self.batch_fmtstr.format(batch)]
entries += [str(meter) for meter in self.meters]
print(' '.join(entries))
def _get_batch_fmtstr(self, num_batches):
num_digits = len(str(num_batches // 1))
fmt = '{:' + str(num_digits) + 'd}'
return '[' + fmt + '/' + fmt.format(num_batches) + ']'
if __name__ == '__main__':
main(args)
Resolved this problem by editing the train_transforms and validate_transforms by adding transforms.Grayscale(3) so that the Grayscale Image shapes go from [1,224,224] to [3,224,224]. Now the Resnet18 has no problem with these images.
train_transforms = transforms.Compose([
transforms.RandomResizedCrop(args.resolution),
transforms.Grayscale(3),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize,
])