在这个pytorch脚本中，我可以在哪里更改广播形状？

zdwk9cvp  于 2023-10-20  发布在  其他
关注(0)|答案(1)|浏览(91)

下面的脚本基本上99%来自NVIDIA Jetson AI github，它使用Pytorch训练不同的DNN进行图像分类。通常，脚本需要命令行参数来提供训练数据的路径。我将其更改为使用Fashion MNIST数据集（第156行f）。
问题是，脚本需要RGB图像，因此输出形状是[3，x，y]，但Fashion MNIST是灰度的，只有一个通道。有没有办法在脚本中改变这一点，或者我需要修改使用的预训练模型？
GIthub with source train.py

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)
pytorch

1条答案

按热度按时间
vlurs2pr

vlurs2pr1#

通过添加transforms.Grayscale（3）编辑train_transforms和validate_transforms来解决此问题，以便灰度图像形状从[1，224，224]变为[3，224，224]。现在Resnet18对这些图像没有问题。

train_transforms = transforms.Compose([
    transforms.RandomResizedCrop(args.resolution),
    transforms.Grayscale(3),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    normalize,
])
赞(0）分享  回复(0）举报  2023-10-20
我来回答

相关问题

    查看更多

    热门标签

    更多
    JavaquerypythonNode开发语言requestUtil数据库Table后端算法LoggerMessageElementParser

    最新问答

    更多
    • 蜀ICP备2022004421号-2 NULL123 https://www.null123.com © All rights reserved
    • 本站内容来源互联网,如果侵犯您的权益请联系我们删除, 联系方式：448109455@qq.com