def evalModel(args, imgs, model): t_test = transforms.Compose([ DataUtil.ToPIL(), # DataUtil.RandomFlips(), DataUtil.TenCrop(args.cropSize, [0.59008044], np.sqrt([0.06342617])), ]) labels = torch.zeros(imgs.shape[0]) test = dataloader.npdataset(imgs, labels.view(-1), t_test) loader = torch.utils.data.DataLoader(test, batch_size=args.bSize, shuffle=False, num_workers=0) outs = [] usegpu = torch.cuda.is_available() if usegpu: model.cuda() # # Iterate. for data in loader: gc.collect() inputs, labels_cpu = data['img'], data['label'] if usegpu: labels_cpu.squeeze_() inputs, labels = Variable(inputs, requires_grad=False).cuda( async=True), Variable(labels_cpu, requires_grad=False).cuda(async=True) else: inputs, labels = Variable(inputs, requires_grad=False), Variable( labels_cpu, requires_grad=False) # # The 5 crop from above takes the corners of the iamge and center. # We must now average the contributions. bs, ncrops, c, h, w = inputs.size() result = model(inputs.view(-1, c, h, w)) # fuse batch size and ncrops out = result.view(bs, ncrops, -1).mean(1) # avg over crops # print(outs) outs.append(copy.copy(out.cpu())) outs = torch.cat(outs, dim=0) return outs
def train(args, imgs, labels, img_val, label_val, modelConst): # # Create the image augmentation. t = transforms.Compose([ DataUtil.ToPIL(), DataUtil.RandomFlips(), # DataUtil.RandomRotation(5), # DataUtil.ColourJitter(0.1, 0.1, 0.1, 0), DataUtil.RandomResizedCrop(args.cropSize, (0.5, 1.3)), DataUtil.ToTensor(), DataUtil.Normalize([0.59008044], np.sqrt([0.06342617])), # DataUtil.TenCrop(140, [0.59008044], np.sqrt([0.06342617])), # ]) t_test = transforms.Compose([ DataUtil.ToPIL(), DataUtil.RandomFlips(), DataUtil.TenCrop(args.cropSize, [0.59008044], np.sqrt([0.06342617])), ]) # RandomRotation # FiveCrop # RandomResizedCrop(size, scale=(0.08, 1.0), ratio=(0.75, 1.3333333333333333), interpolation=2) topil = transforms.ToPILImage() train = dataloader.npdataset(imgs, labels.view(-1), t) validation = dataloader.npdataset(img_val, label_val.view(-1), t_test) stages = { 'train': torch.utils.data.DataLoader(train, batch_size=args.bSize, shuffle=True, num_workers=0, pin_memory=True), 'val': torch.utils.data.DataLoader(validation, batch_size=args.bSize, shuffle=False, num_workers=0, pin_memory=True), } model = modelConst() usegpu = torch.cuda.is_available() criteria = nn.CrossEntropyLoss() # # Whether to use the GPU. if usegpu: model.cuda() # # Type of optimizer. optimizer = torch.optim.Adam(model.parameters(), lr=0.0001, betas=(0.9, 0.999), eps=1e-08, weight_decay=1e-2) bestModel = copy.deepcopy(model.state_dict()) bestAcc = 0 logger = None logEpoch = doNothing closeLogger = doNothing if args.useTB: logger = SummaryWriter() logEpoch = logEpochTensorboard closeLogger = closeTensorboard # # Iterate. for epoch in range(args.numEpochs): printColour('Epoch {}/{}'.format(epoch, args.numEpochs - 1), colours.OKBLUE) for stage in stages: gc.collect() print('Stage: ', stage) # # Switch on / off gradients. model.train(stage == 'train') # # The current loss. runningLoss = 0.0 runningCorrect = 0.0 loader = stages[stage] # # Progress bar. numMini = len(loader) pbar = progressbar.ProgressBar(max_value=numMini - 1) # # Train. for i, data in enumerate(loader): inputs_cpu, labels_cpu = data['img'], data['label'] if usegpu: labels_cpu.squeeze_() inputs, labels = Variable( inputs_cpu, requires_grad=False).cuda(async=True), Variable( labels_cpu, requires_grad=False).cuda(async=True) else: inputs, labels = Variable(inputs_cpu, requires_grad=False, volatile=True), Variable( labels_cpu, requires_grad=False, volatile=True) # # Forward through network. if stage == 'train': out = model(inputs) else: # # The 5 crop from above takes the corners of the iamge and center. # We must now average the contributions. bs, ncrops, c, h, w = inputs_cpu.size() inputs = inputs.view(-1, c, h, w) result = model(inputs) # fuse batch size and ncrops out = result.view(bs, ncrops, -1).mean(1) # avg over crops # # Backward pass. optimizer.zero_grad() _, preds = torch.max(out.data, 1) loss = criteria(out, labels) # # Backwards pass. if stage == 'train': loss.backward() optimizer.step() dCorrect = torch.sum(preds == labels.data) # # Stats. runningLoss += loss.data[0] runningCorrect += dCorrect pbar.update(i) # # Overall stats epochLoss = runningLoss / len(stages[stage]) epochAcc = runningCorrect / (len(stages[stage]) * args.bSize) # # Check if we have the new best model. isBest = False if stage == 'val' and epochAcc > bestAcc: isBest = True bestAcc = epochAcc bestModel = model.state_dict() # # Print per epoch results. print('\n{} Loss: {:.4f} Acc: {:.4f}'.format( stage, epochLoss, epochAcc)) # # Summary for logging in TB. summary = { 'phase': stage, 'epoch': epoch, 'loss': epochLoss, 'acc': epochAcc, 'data': data, 'pred': preds } logEpoch(logger, model, summary) printColour('Best validation performance:%f' % (bestAcc), colours.OKGREEN) closeLogger(logger) retModel = copy.deepcopy(bestModel) for key, val in retModel.items(): retModel[key] = val.cpu() return retModel, bestAcc