def main(): args = parser.parse_args() if args.num_classes == 1000: model = alexnet(pretrained=True).to("cuda:0") print("Load the official pre-training weight successfully.") else: model = AlexNet(num_classes=args.num_classes).to("cuda:0") model.load_state_dict( torch.load("/opt/tensorrt_models/torch/alexnet/alexnet.pth")) print("Load the specified pre-training weight successfully.") model.eval() f = open("/opt/tensorrt_models/torch/alexnet/alexnet.wts", "w") f.write("{}\n".format(len(model.state_dict().keys()))) for k, v in model.state_dict().items(): vr = v.reshape(-1).cpu().numpy() f.write("{} {}".format(k, len(vr))) for vv in vr: f.write(" ") f.write(struct.pack(">f", float(vv)).hex()) f.write("\n") print( "The weight conversion has been completed and saved to `/opt/tensorrt_models/torch/alexnet/alexnet.wts`." )
def predictions_AlexNet(imagePath): # Open image input_image = Image.open(imagePath) # Preprocess image preprocess = transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) input_tensor = preprocess(input_image) input_batch = input_tensor.unsqueeze(0) # create a mini-batch as expected by the model # Load class names labels_map = json.load(open("AlexNet_labels_map.txt")) labels_map = [labels_map[str(i)] for i in range(1000)] # Classify with AlexNet model = AlexNet.from_pretrained("alexnet") model.eval() with torch.no_grad(): logits = model(input_batch) preds = torch.topk(logits, k=5).indices.squeeze(0).tolist() res_list = [] for idx in preds: label = labels_map[idx] prob = torch.softmax(logits, dim=1)[0, idx].item() res_list.append((label, prob)) return res_list
# ============================================================================== import json import os import urllib.request import torch import torchvision.transforms as transforms from PIL import Image from django.shortcuts import render from rest_framework.views import APIView from alexnet_pytorch import AlexNet device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") model = AlexNet.from_pretrained("alexnet") # move the model to GPU for speed if available model = model.to(device) # switch to evaluate mode model.eval() def preprocess(filename, label): input_image = Image.open(filename) transform = transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
def main_worker(gpu, ngpus_per_node, args): global best_acc1 args.gpu = gpu if args.gpu is not None: print(f"Use GPU: {args.gpu} for training!") if args.distributed: if args.dist_url == "env://" and args.rank == -1: args.rank = int(os.environ["RANK"]) if args.multiprocessing_distributed: # For multiprocessing distributed training, rank needs to be the # global rank among all the processes args.rank = args.rank * ngpus_per_node + gpu dist.init_process_group(backend=args.dist_backend, init_method=args.dist_url, world_size=args.world_size, rank=args.rank) # create model if 'alexnet' in args.arch: if args.pretrained: model = AlexNet.from_pretrained(args.arch, args.num_classes) print(f"=> using pre-trained model '{args.arch}'") else: print(f"=> creating model '{args.arch}'") model = AlexNet.from_name(args.arch) else: 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]() if args.distributed: # For multiprocessing distributed, DistributedDataParallel constructor # should always set the single device scope, otherwise, # DistributedDataParallel will use all available devices. if args.gpu is not None: torch.cuda.set_device(args.gpu) model.cuda(args.gpu) # When using a single GPU per process and per # DistributedDataParallel, we need to divide the batch size # ourselves based on the total number of GPUs we have args.batch_size = int(args.batch_size / ngpus_per_node) args.workers = int(args.workers / ngpus_per_node) model = torch.nn.parallel.DistributedDataParallel( model, device_ids=[args.gpu]) else: model.cuda() # DistributedDataParallel will divide and allocate batch_size to all # available GPUs if device_ids are not set model = torch.nn.parallel.DistributedDataParallel(model) elif args.gpu is not None: torch.cuda.set_device(args.gpu) model = model.cuda(args.gpu) else: # DataParallel will divide and allocate batch_size to all available # GPUs if args.arch.startswith('alexnet'): model.features = torch.nn.DataParallel(model.features) model.cuda() else: model = torch.nn.DataParallel(model).cuda() # define loss function (criterion) and optimizer criterion = nn.CrossEntropyLoss().cuda(args.gpu) optimizer = torch.optim.SGD(model.parameters(), args.lr, momentum=args.momentum, weight_decay=args.weight_decay) if mixed_precision: model, optimizer = amp.initialize(model, optimizer, opt_level=args.opt_level) # 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'] best_acc1 = checkpoint['best_acc1'] if args.gpu is not None: # best_acc1 may be from a checkpoint from a different GPU best_acc1 = best_acc1.to(args.gpu) model.load_state_dict(checkpoint['state_dict']) optimizer.load_state_dict(checkpoint['optimizer']) amp.load_state_dict(checkpoint['amp']) print(f"=> loaded checkpoint '{args.resume}' (epoch {checkpoint['epoch']})") else: print(f"=> no checkpoint found at '{args.resume}'") cudnn.benchmark = True # Data loading code normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) train_dataset = datasets.CIFAR100( root=args.data, download=True, train=True, transform=transforms.Compose([ transforms.Resize(224), transforms.RandomHorizontalFlip(), transforms.ToTensor(), normalize, ])) if args.distributed: train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset) else: train_sampler = None train_loader = torch.utils.data.DataLoader( train_dataset, batch_size=args.batch_size, shuffle=(train_sampler is None), num_workers=args.workers, pin_memory=True, sampler=train_sampler) val_dataset = datasets.CIFAR100( root=args.data, download=True, train=False, transform=transforms.Compose([ transforms.Resize(224), transforms.ToTensor(), normalize, ])) val_loader = torch.utils.data.DataLoader( val_dataset, batch_size=args.batch_size, shuffle=False, num_workers=args.workers, pin_memory=True) if args.evaluate: top1 = validate(val_loader, model, criterion, args) with open('res.txt', 'w') as f: print(f"Acc@1: {top1}", file=f) return for epoch in range(args.start_epoch, args.epochs): if args.distributed: train_sampler.set_epoch(epoch) adjust_learning_rate(optimizer, epoch, args) # train for one epoch train(train_loader, model, criterion, optimizer, epoch, args) # evaluate on validation set acc1 = validate(val_loader, model, criterion, args) # remember best acc@1 and save checkpoint is_best = acc1 > best_acc1 best_acc1 = max(acc1, best_acc1) if not args.multiprocessing_distributed or (args.multiprocessing_distributed and args.rank % ngpus_per_node == 0): save_checkpoint({ 'epoch': epoch + 1, 'arch': args.arch, 'state_dict': model.state_dict(), 'best_acc1': best_acc1, 'optimizer': optimizer.state_dict(), 'amp': amp.state_dict(), }, is_best)
def main_worker(gpu, ngpus_per_node, args): global best_acc1 args.gpu = gpu if args.gpu is not None: print(f"Use GPU: {args.gpu} for training!") if args.distributed: if args.dist_url == "env://" and args.rank == -1: args.rank = int(os.environ["RANK"]) if args.multiprocessing_distributed: # For multiprocessing distributed training, rank needs to be the # global rank among all the processes args.rank = args.rank * ngpus_per_node + gpu dist.init_process_group(backend=args.dist_backend, init_method=args.dist_url, world_size=args.world_size, rank=args.rank) # create model if args.pretrained and args.num_classes == 1000: model = alexnet(pretrained=True) print(f"=> using pre-trained model") elif args.pretrained: print(f"=> creating model and Fine-tuning pre-trained model.") model = alexnet(pretrained=True) model.classifier = nn.Sequential( nn.Dropout(), nn.Linear(256 * 6 * 6, 4096), nn.ReLU(inplace=True), nn.Dropout(), nn.Linear(4096, 4096), nn.ReLU(inplace=True), nn.Linear(4096, args.num_classes), ) args.lr = 0.01 else: print(f"=> creating model.") model = AlexNet(num_classes=args.num_classes) if args.distributed: # For multiprocessing distributed, DistributedDataParallel constructor # should always set the single device scope, otherwise, # DistributedDataParallel will use all available devices. if args.gpu is not None: torch.cuda.set_device(args.gpu) model.cuda(args.gpu) # When using a single GPU per process and per # DistributedDataParallel, we need to divide the batch size # ourselves based on the total number of GPUs we have args.batch_size = int(args.batch_size / ngpus_per_node) args.workers = int(args.workers / ngpus_per_node) model = torch.nn.parallel.DistributedDataParallel( model, device_ids=[args.gpu]) else: model.cuda() # DistributedDataParallel will divide and allocate batch_size to all # available GPUs if device_ids are not set model = torch.nn.parallel.DistributedDataParallel(model) elif args.gpu is not None: torch.cuda.set_device(args.gpu) model = model.cuda(args.gpu) else: # DataParallel will divide and allocate batch_size to all available # GPUs model.features = torch.nn.DataParallel(model.features) model.cuda() # define loss function (criterion) and optimizer criterion = nn.CrossEntropyLoss().cuda(args.gpu) optimizer = torch.optim.SGD(model.parameters(), args.lr, momentum=args.momentum, weight_decay=args.weight_decay) # 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"] best_acc1 = checkpoint["best_acc1"] if args.gpu is not None: # best_acc1 may be from a checkpoint from a different GPU best_acc1 = best_acc1.to(args.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 # Data loading code traindir = os.path.join(args.data, "train") valdir = os.path.join(args.data, "val") normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) train_dataset = datasets.ImageFolder(traindir, transform=transforms.Compose([ transforms.RandomResizedCrop( args.image_size), transforms.RandomHorizontalFlip(), transforms.ToTensor(), normalize, ])) if args.distributed: train_sampler = torch.utils.data.distributed.DistributedSampler( train_dataset) else: train_sampler = None train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=args.batch_size, shuffle=(train_sampler is None), num_workers=args.workers, pin_memory=True, sampler=train_sampler) val_transforms = transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), normalize, ]) print("Using image size", 224) val_loader = torch.utils.data.DataLoader(datasets.ImageFolder( valdir, val_transforms), batch_size=args.batch_size, shuffle=False, num_workers=args.workers, pin_memory=True) if args.evaluate: top1 = validate(val_loader, model, criterion, args) with open("res.txt", "w") as f: print(f"Acc@1: {top1}", file=f) return for epoch in range(args.start_epoch, args.epochs): if args.distributed: train_sampler.set_epoch(epoch) adjust_learning_rate(optimizer, epoch, args) # train for one epoch train(train_loader, model, criterion, optimizer, epoch, args) # evaluate on validation set acc1 = validate(val_loader, model, criterion, args) # remember best acc@1 and save checkpoint is_best = acc1 > best_acc1 best_acc1 = max(acc1, best_acc1) if not args.multiprocessing_distributed or ( args.multiprocessing_distributed and args.rank % ngpus_per_node == 0): save_checkpoint( { "epoch": epoch + 1, "state_dict": model.state_dict(), "best_acc1": best_acc1, "optimizer": optimizer.state_dict(), }, is_best)
def openimage(self): img_name, _ = QFileDialog.getOpenFileName(self, "Open image", "", "*.jpg;;*.png;;All Files(*)") jpg = QtGui.QPixmap(img_name).scaled(args.image_size, args.image_size) self.label.setPixmap(jpg) text, prob = classifier(img_name) self.echo(str(text), prob) def echo(self, text, prob): QMessageBox.information( self, "Message", f"Label :{str(text):<75}\nProbability: {prob:.6f}") if __name__ == "__main__": model = AlexNet.from_pretrained(args.model_name) model.eval() if args.image_size is None: args.image_size = AlexNet.get_image_size(args.model_name) # Preprocess image tfms = transforms.Compose([ transforms.Resize(args.image_size), transforms.CenterCrop(args.image_size), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), ]) # Load class names labels_map = json.load(open(args.labels_map)) labels_map = [labels_map[str(i)] for i in range(args.num_classes)]
def __init__(self, model=None): if (model is None): self.alex = AlexNet.from_pretrained('alexnet', num_classes=2) else: self.alex = model