class FaceGenderAge(object): WRN_WEIGHTS_PATH = "https://github.com/Tony607/Keras_age_gender/releases/download/V1.0/weights.18-4.06.hdf5" def __init__(self, depth=16, width=8, face_size=64): self.face_size = face_size self.model = WideResNet(face_size, depth=depth, k=width)() model_dir = os.path.join(os.getcwd(), "models") fpath = get_file('weights.18-4.06.hdf5', self.WRN_WEIGHTS_PATH, cache_subdir=model_dir) self.model.load_weights(fpath) self.model._make_predict_function() def predict_agge(self, frame): inference_frame = cv2.resize(frame, (self.face_size, self.face_size)) inference_frame = inference_frame[np.newaxis, :, :, :] results = self.model.predict(inference_frame) predicted_genders = results[0] ages = np.arange(0, 101).reshape(101, 1) predicted_ages = results[1].dot(ages).flatten() label = "Age: {}, Gender: {}".format( int(predicted_ages), "F" if predicted_genders[0][0] > 0.5 else "M") return label
def load_network(loc): net_checkpoint = torch.load(loc) start_epoch = net_checkpoint['epoch'] SavedConv, SavedBlock = what_conv_block(net_checkpoint['conv'], net_checkpoint['blocktype'], net_checkpoint['module']) net = WideResNet(args.wrn_depth, args.wrn_width, SavedConv, SavedBlock, num_classes=num_classes, dropRate=0).cuda() net.load_state_dict(net_checkpoint['net']) return net, start_epoch
def __init__(self, depth=16, width=8, face_size=64): self.face_size = face_size self.model = WideResNet(face_size, depth=depth, k=width)() model_dir = os.path.join(os.getcwd(), "models") fpath = get_file('weights.18-4.06.hdf5', self.WRN_WEIGHTS_PATH, cache_subdir=model_dir) self.model.load_weights(fpath) self.model._make_predict_function()
def __init__(self): if torch.cuda.is_available(): self.device = torch.device('cuda') else: self.device = torch.device('cpu') print('WARNING: Found no valid GPU device - Running on CPU') self.model = WideResNet(DEPTH, cfg.NUM_TRANS) self.model.cuda(self.device) self.criterion = torch.nn.CrossEntropyLoss().cuda(self.device) self.optimizer = None
def build_network(Conv, Block, network): if network == 'WideResNet': return WideResNet(28, 10, Conv, Block, num_classes=10, dropRate=0) elif network == 'WRN_50_2': return WRN_50_2(Conv) elif network == 'DARTS': return DARTS(Conv, num_classes=10, drop_path_prob=0., auxiliary=False)
def load_network(loc, masked=False): net_checkpoint = torch.load(loc) start_epoch = net_checkpoint['epoch'] SavedConv, SavedBlock = what_conv_block(net_checkpoint['conv'], net_checkpoint['blocktype'], net_checkpoint['module']) net = WideResNet(args.wrn_depth, args.wrn_width, SavedConv, SavedBlock, num_classes=num_classes, dropRate=0, masked=masked).cuda() if masked: new_sd = net.state_dict() old_sd = net_checkpoint['net'] new_names = [v for v in new_sd] old_names = [v for v in old_sd] for i, j in enumerate(new_names): new_sd[j] = old_sd[old_names[i]] net.load_state_dict(new_sd) else: net.load_state_dict(net_checkpoint['net']) return net, start_epoch
def load_baseline_model(self): """ Load a simple baseline model AND dataset Note that this sets the model to training mode """ if self.args.dataset == DATASET_CIFAR_10: imsize, in_channel, num_classes = 32, 3, 10 elif self.args.dataset == DATASET_CIFAR_100: imsize, in_channel, num_classes = 32, 3, 100 elif self.args.dataset == DATASET_MNIST: imsize, in_channel, num_classes = 28, 1, 10 elif self.args.dataset == DATASET_BOSTON: imsize, in_channel, num_classes = 13, 1, 1 # init_l2 = -7 # TODO: Important to make sure this is small enough to be unregularized when starting? if self.args.model == MODEL_RESNET18: cnn = ResNet18(num_classes=num_classes) elif self.args.model == MODEL_WIDERESNET: cnn = WideResNet(depth=28, num_classes=num_classes, widen_factor=10, dropRate=0.3) elif self.args.model[:3] == MODEL_MLP: cnn = Net(self.args.num_layers, 0.0, imsize, in_channel, INIT_L2, num_classes=num_classes, do_classification=self.args.do_classification) elif self.args.model == MODEL_CNN_MLP: cnn = CNN_MLP(learning_rate=0.0001) checkpoint = None if self.args.load_baseline_checkpoint: checkpoint = torch.load(self.args.load_baseline_checkpoint) cnn.load_state_dict(checkpoint['model_state_dict']) model = cnn.to(self.device) if self.args.use_weight_decay: if self.args.weight_decay_all: num_p = sum(p.numel() for p in model.parameters()) weights = np.ones(num_p) * INIT_L2 model.weight_decay = Variable(torch.FloatTensor(weights).to( self.device), requires_grad=True) else: weights = INIT_L2 model.weight_decay = Variable(torch.FloatTensor([weights]).to( self.device), requires_grad=True) model.weight_decay = model.weight_decay.to(self.device) model.train() return model, checkpoint
def load_baseline_model(args): """ :param args: :return: """ if args.dataset == 'cifar10': num_classes = 10 train_loader, val_loader, test_loader = data_loaders.load_cifar10(args.batch_size, val_split=True, augmentation=args.data_augmentation) elif args.dataset == 'cifar100': num_classes = 100 train_loader, val_loader, test_loader = data_loaders.load_cifar100(args.batch_size, val_split=True, augmentation=args.data_augmentation) elif args.dataset == 'mnist': args.datasize, args.valsize, args.testsize = 100, 100, 100 num_train = args.datasize if args.datasize == -1: num_train = 50000 from data_loaders import load_mnist train_loader, val_loader, test_loader = load_mnist(args.batch_size, subset=[args.datasize, args.valsize, args.testsize], num_train=num_train) if args.model == 'resnet18': cnn = ResNet18(num_classes=num_classes) elif args.model == 'wideresnet': cnn = WideResNet(depth=28, num_classes=num_classes, widen_factor=10, dropRate=0.3) checkpoint = None if args.load_baseline_checkpoint: checkpoint = torch.load(args.load_baseline_checkpoint) cnn.load_state_dict(checkpoint['model_state_dict']) model = cnn.cuda() model.train() return model, train_loader, val_loader, test_loader, checkpoint
def build_network(Conv, Block): if args.network == 'WideResNet': return WideResNet(args.wrn_depth, args.wrn_width, Conv, Block, num_classes=num_classes, dropRate=0, s=args.AT_split) elif args.network == 'WRN_50_2': return WRN_50_2(Conv) elif args.network == 'MobileNetV2': return MobileNetV2(Conv) elif args.network == 'DARTS': return DARTS(Conv, num_classes=num_classes)
def build_network(Conv, Block): if args.network == 'WideResNet': return WideResNet(args.wrn_depth, args.wrn_width, Conv, Block, num_classes=num_classes, dropRate=0, s=args.AT_split, spectral=args.spectral and not (args.rank_scale or args.target_ratio)) elif args.network == 'WRN_50_2': return WRN_50_2(Conv) elif args.network == 'MobileNetV2': return MobileNetV2(Conv) elif args.network == 'DARTS': return DARTS(Conv, num_classes=num_classes)
def random_conv_sub(net, param_target): conv_options = get_ordered_conv_options() while ((get_no_params(net) > param_target) or (get_no_params(net) < (param_target - (param_target * 0.05)))): convs = [] for m in net.modules(): conv = choice(conv_options) convs.append(conv) net = WideResNet(args.wrn_depth, args.wrn_width, Conv, MaskBlock, num_classes=num_classes, dropRate=0, s=args.AT_split, convs=convs).cuda() return convs
def build_network(Conv, Block): if args.network == 'WideResNet': return WideResNet(args.wrn_depth, args.wrn_width, Conv, Block, num_classes=num_classes, dropRate=0) elif args.network == 'WRN_50_2': return WRN_50_2(Conv) elif args.network == 'DARTS': assert not args.conv == 'Conv', 'The base network here used' \ ' separable convolutions, so you probably did not mean to set this' \ ' option.' return DARTS(Conv, num_classes=num_classes, drop_path_prob=0., auxiliary=False) elif args.network == 'MobileNetV2': return MobileNetV2(Conv)
total_nonzero = 0.0 # With the valid epsilon, we can set sparsities of the remaning layers. for name, mask in masks.items(): n_param = np.prod(mask.shape) if name in dense_layers: density_dict[name] = 1.0 else: probability_one = epsilon * raw_probabilities[name] density_dict[name] = probability_one logging.info( f"layer: {name}, shape: {mask.shape}, density: {density_dict[name]}" ) total_nonzero += density_dict[name] * mask.numel() logging.info(f"Overall sparsity {total_nonzero/total_params}") return density_dict if __name__ == "__main__": from models.wide_resnet import WideResNet model = WideResNet(depth=22, widen_factor=2) logging.basicConfig() logging.getLogger().setLevel(logging.INFO) tim_ERK(model, density=0.2) logging.info("========") googleAI_ERK(model, density=0.2)
emotion_labels = get_labels('fer2013') gender_labels = get_labels('imdb') race_labels = get_labels('race') font = cv2.FONT_HERSHEY_SIMPLEX # hyper-parameters for bounding boxes shape frame_window = 10 gender_offsets = (30, 60) emotion_offsets = (20, 40) # loading models face_detection = load_detection_model(detection_model_path) emotion_classifier = load_model(emotion_model_path, compile=False) # race_classifier = load_model(race_model_path, compile=False) age_gender_classifier = WideResNet(64, depth=16, k=8)() age_gender_classifier.load_weights(age_gender_model_path) # getting input model shapes for inference emotion_target_size = emotion_classifier.input_shape[1:3] age_gender_target_size = age_gender_classifier.input_shape[1:3] # race_target_size = race_classifier.input_shape[1:3] # starting lists for calculating modes gender_window = [] emotion_window = [] # starting video streaming cv2.namedWindow('window_frame') video_capture = cv2.VideoCapture(video_path)
test_transform = transforms.Compose([transforms.ToTensor(), normalize]) if args.dataset == 'cifar10': num_classes = 10 train_loader, val_loader, test_loader = data_loaders.load_cifar10( args.batch_size, val_split=True, augmentation=args.data_augmentation) elif args.dataset == 'cifar100': num_classes = 100 train_loader, val_loader, test_loader = data_loaders.load_cifar100( args.batch_size, val_split=True, augmentation=args.data_augmentation) if args.model == 'resnet18': cnn = ResNet18(num_classes=num_classes) elif args.model == 'wideresnet': cnn = WideResNet(depth=28, num_classes=num_classes, widen_factor=10, dropRate=0.3) cnn = cnn.cuda() criterion = nn.CrossEntropyLoss().cuda() if args.optimizer == 'sgdm': cnn_optimizer = torch.optim.SGD(cnn.parameters(), lr=args.lr, momentum=0.9, nesterov=True, weight_decay=args.wdecay) elif args.optimizer == 'sgd': cnn_optimizer = torch.optim.SGD(cnn.parameters(), lr=args.lr, weight_decay=args.wdecay)
class NeuralNet: def __init__(self): if torch.cuda.is_available(): self.device = torch.device('cuda') else: self.device = torch.device('cpu') print('WARNING: Found no valid GPU device - Running on CPU') self.model = WideResNet(DEPTH, cfg.NUM_TRANS) self.model.cuda(self.device) self.criterion = torch.nn.CrossEntropyLoss().cuda(self.device) self.optimizer = None def test(self, test_gen): self.model.eval() batch_time = self.AverageMeter('Time', ':6.3f') losses = self.AverageMeter('Loss', ':.4e') top1 = self.AverageMeter('Acc@1', ':6.2f') progress = self.ProgressMeter(len(test_gen), batch_time, losses, top1, prefix='Test: ') # switch to evaluate mode self.model.eval() with torch.no_grad(): end = time.time() for i, (input, target, _) in enumerate(test_gen): input = input.cuda(self.device, non_blocking=True) target = target.cuda(self.device, non_blocking=True) # Compute output output, _ = self.model([input, target]) loss = self.criterion(output, target) # Measure accuracy and record loss acc1 = self._accuracy(output, target, topk=(1)) losses.update(loss.item(), input.size(0)) top1.update(acc1[0].cpu().detach().item(), input.size(0)) # Measure elapsed time batch_time.update(time.time() - end) end = time.time() # Print to screen if i % 100 == 0: progress.print(i) # TODO: this should also be done with the ProgressMeter print(' * Acc@1 {top1.avg:.3f}'.format(top1=top1)) return top1.avg def train(self, train_gen, test_gen, epochs, lr=0.0001, lr_plan=None, momentum=0.9, wd=5e-4): self.optimizer = torch.optim.SGD(self.model.parameters(), lr=lr, momentum=momentum, weight_decay=wd) #self.optimizer = torch.optim.Adam(self.model.parameters()) for epoch in range(epochs): self._adjust_lr_rate(self.optimizer, epoch, lr_plan) print("=> Training (specific label)") self._train_step(train_gen, epoch, self.optimizer) print("=> Validation (entire dataset)") self.test(test_gen) def _train_step(self, train_gen, epoch, optimizer): self.model.train() batch_time = self.AverageMeter('Time', ':6.3f') data_time = self.AverageMeter('Data', ':6.3f') losses = self.AverageMeter('Loss', ':.4e') top1 = self.AverageMeter('Acc@1', ':6.2f') progress = self.ProgressMeter(len(train_gen), batch_time, data_time, losses, top1, prefix="Epoch: [{}]".format(epoch)) end = time.time() for i, (input, target, _) in enumerate(train_gen): # measure data loading time data_time.update(time.time() - end) input = input.cuda(self.device, non_blocking=True) target = target.cuda(self.device, non_blocking=True) # Compute output output, trans_out = self.model([input, target]) loss = self.criterion(output, target) # measure accuracy and record loss acc1 = self._accuracy(output, target, topk=(1)) losses.update(loss.item(), input.size(0)) top1.update(acc1[0].cpu().detach().item(), input.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 % 100 == 0: progress.print(i) def evaluate(self, eval_gen): # switch to evaluate mode self.model.eval() score_func_list = [] labels_list = [] with torch.no_grad(): for i, (input, target, labels) in enumerate(eval_gen): input = input.cuda(self.device, non_blocking=True) #Target- the transforamation class target = target.cuda(self.device, non_blocking=True) #The true label labels = labels[[ cfg.NUM_TRANS * x for x in range(len(labels) // cfg.NUM_TRANS) ]].cuda(self.device, non_blocking=True) # Compute output # #TODO: Rewrite this code section, can be more efficient output_SM = self.model([input, target]) target_mat = torch.zeros_like(output_SM[0]) target_mat[range(output_SM[0].shape[0]), target] = 1 target_SM = (target_mat * output_SM[0]).sum(dim=1).view( -1, cfg.NUM_TRANS).sum(dim=1) score_func_list.append(1 / cfg.NUM_TRANS * target_SM) labels_list.append(labels) return torch.cat(score_func_list), torch.cat(labels_list) def _adjust_lr_rate(self, optimizer, epoch, lr_dict): if lr_dict is None: return for key, value in lr_dict.items(): if epoch == key: print("=> New learning rate set of {}".format(value)) for param_group in optimizer.param_groups: param_group['lr'] = value def summary(self, x_size, print_it=True): return self.model.summary(x_size, print_it=print_it) def print_weights(self): self.model.print_weights() @staticmethod def _accuracy(output, target, topk=(1)): """Computes the accuracy over the k top predictions for the specified values of k""" with torch.no_grad(): maxk = 1 batch_size = target.size(0) _, pred = output.topk(maxk, 1, True, True) pred = pred.t() correct = pred.eq(target.view(1, -1).expand_as(pred)) res = [] correct_k = correct[0].view(-1).float().sum(0, keepdim=True) res.append(correct_k.mul_(100.0 / batch_size)) return res 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): def __init__(self, num_batches, *meters, prefix=""): self.batch_fmtstr = self._get_batch_fmtstr(num_batches) self.meters = meters self.prefix = prefix def print(self, batch): entries = [self.prefix + self.batch_fmtstr.format(batch)] entries += [str(meter) for meter in self.meters] print('\t'.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) + ']'
num_workers=4) test_set = torchvision.datasets.FashionMNIST(root="/mnt/ds/ryo", train=False, download=True, transform=transform_test) test_loader = torch.utils.data.DataLoader(test_set, batch_size=100, shuffle=False, num_workers=4) num_class = 10 ############## parameter ################### if str(args.model) == "resnet18": model = ResNet18(num_class).to(device) if str(args.model) == "wide-resnet28-10": model = WideResNet(28, 10, 0, in_channels=3, labels=num_class).to(device) # optimizing if str(args.optimizer) == "SGD": optimizer = optim.SGD(model.parameters(), lr=LEANRATE, momentum=0) if str(args.optimizer) == "SAM": base_optimizer = torch.optim.SGD optimizer = SAM(model.parameters(), base_optimizer, rho=0.05, lr=LEANRATE, momentum=0, weight_decay=0.0005) criterion = nn.CrossEntropyLoss()
net_checkpoint = torch.load(loc) start_epoch = net_checkpoint['epoch'] SavedConv, SavedBlock = what_conv_block(net_checkpoint['conv'], net_checkpoint['blocktype'], net_checkpoint['module']) net = WideResNet(args.wrn_depth, args.wrn_width, SavedConv, SavedBlock, num_classes=num_classes, dropRate=0).cuda() net.load_state_dict(net_checkpoint['net']) return net, start_epoch if args.mode == 'teacher': if args.resume: print('Mode Teacher: Loading teacher and continuing training...') teach, start_epoch = load_network('checkpoints/%s.t7' % args.teacher_checkpoint) else: print('Mode Teacher: Making a teacher network from scratch and training it...') teach = WideResNet(args.wrn_depth, args.wrn_width, Conv, Block, num_classes=num_classes, dropRate=0).cuda() get_no_params(teach) optimizer = optim.SGD(teach.parameters(), lr=args.lr, momentum=0.9, weight_decay=args.weightDecay) scheduler = lr_scheduler.MultiStepLR(optimizer, milestones=epoch_step, gamma=args.lr_decay_ratio) # Decay the learning rate depending on the epoch for e in range(0,start_epoch): scheduler.step() for epoch in tqdm(range(start_epoch, args.epochs)): scheduler.step() print('Teacher Epoch %d:' % epoch) print('Learning rate is %s' % [v['lr'] for v in optimizer.param_groups][0]) writer.add_scalar('learning_rate', [v['lr'] for v in optimizer.param_groups][0], epoch)
net.load_state_dict(net_checkpoint['net']) return net, start_epoch if args.mode == 'teacher': if args.resume: print('Mode Teacher: Loading teacher and continuing training...') teach, start_epoch = load_network('checkpoints/%s.t7' % args.teacher_checkpoint) else: print( 'Mode Teacher: Making a teacher network from scratch and training it...' ) teach = WideResNet(args.wrn_depth, args.wrn_width, Conv, Block, num_classes=num_classes, dropRate=0).cuda() get_no_params(teach) optimizer = optim.SGD(teach.parameters(), lr=args.lr, momentum=0.9, weight_decay=args.weightDecay) scheduler = lr_scheduler.MultiStepLR(optimizer, milestones=epoch_step, gamma=args.lr_decay_ratio) # Decay the learning rate depending on the epoch for e in range(0, start_epoch): scheduler.step()
def detection(self): depth = 16 width = 8 img_size = 64 model = WideResNet(img_size, depth=depth, k=width)() model.load_weights(r'models/weights.hdf5') detector = dlib.get_frontal_face_detector() image_np = cv2.imdecode(np.fromfile(self.fname, dtype=np.uint8), -1) img_h = image_np.shape[0] img_w = image_np.shape[1] detected = detector(image_np, 1) gender_faces = [] labels = [] original_faces = [] photo_position = [] change_male_to_female_path = r'models/netG_A2B.pth' change_female_to_male_path = r'models/netG_B2A.pth' # 加载CycleGAN模型 netG_male2female = Generator(3, 3) netG_female2male = Generator(3, 3) netG_male2female.load_state_dict( torch.load(change_male_to_female_path, map_location='cpu')) netG_female2male.load_state_dict( torch.load(change_female_to_male_path, map_location='cpu')) # 设置模型为预测模式 netG_male2female.eval() netG_female2male.eval() transform_list = [ transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) ] transform = transforms.Compose(transform_list) """ 这段内容为图片数据处理frame """ if len(detected) > 0: for i, d in enumerate(detected): # weight和height表示原始图片的宽度和高度,因为我们需要不停地resize处理图片 # 最后要贴回到原始图片中去,w和h就用来做最后的resize x0, y0, x1, y1, w, h = d.left(), d.top(), d.right(), d.bottom( ), d.width(), d.height() x0 = max(int(x0 - 0.25 * w), 0) y0 = max(int(y0 - 0.45 * h), 0) x1 = min(int(x1 + 0.25 * w), img_w - 1) y1 = min(int(y1 + 0.05 * h), img_h - 1) w = x1 - x0 h = y1 - y0 if w > h: x0 = x0 + w // 2 - h // 2 w = h x1 = x0 + w else: y0 = y0 + h // 2 - w // 2 h = w y1 = y0 + h original_faces.append( cv2.resize(image_np[y0:y1, x0:x1, :], (256, 256))) gender_faces.append( cv2.resize(image_np[y0:y1, x0:x1, :], (img_size, img_size))) photo_position.append([y0, y1, x0, x1, w, h]) gender_faces = np.array(gender_faces) results = model.predict(gender_faces) predicted_genders = results[0] for i in range(len(original_faces)): labels.append('F' if predicted_genders[i][0] > 0.5 else 'M') for i, gender in enumerate(labels): # 这几个变量用于接下来图片缩放和替换 y0, y1, x0, x1, w, h = photo_position[i] # 将数据转换成Pytorch可以处理的格式 picture = transform(original_faces[i]) picture = Variable(picture) input_picture = picture.unsqueeze(0) if gender == "M": fake_female = 0.5 * (netG_male2female(input_picture).data + 1.0) out_img = fake_female.detach().squeeze(0) else: fake_male = 0.5 * (netG_female2male(input_picture).data + 1.0) out_img = fake_male.detach().squeeze(0) # 需要将Pytorch处理之后得到的数据,转换为OpenCV可以处理的格式 # 下面代码就是转换代码 image_numpy = out_img.float().numpy() image_numpy = (np.transpose(image_numpy, (1, 2, 0)) + 1) / 2.0 * 255.0 image_numpy = image_numpy.clip(0, 255) image_numpy = image_numpy.astype(np.uint8) # 将转换好的性别图片替换到原始图片中去 # 使用泊松融合使生成图像和背景图像浑然一体 # 使用方法:cv2.seamlessClone(src, dst, mask, center, flags) generate_face = cv2.resize(image_numpy, (w, h)) # Create an all white mask, 感兴趣的需要替换的目标区域,精确地mask可以更好的替换,这里mask就是生成图片的大小 mask = 255 * np.ones((w, h), image_np.dtype) # center是目标影像的中心在背景图像上的坐标! center_y = y0 + h // 2 center_x = x0 + w // 2 center = (center_x, center_y) output_face = cv2.seamlessClone(generate_face, image_np, mask, center, cv2.NORMAL_CLONE) self.out_img = output_face
def main(args): harakiri = Harakiri() harakiri.set_max_plateau(20) train_loss_meter = Meter() val_loss_meter = Meter() val_accuracy_meter = Meter() log = JsonLogger(args.log_path, rand_folder=True) log.update(args.__dict__) state = args.__dict__ state['exp_dir'] = os.path.dirname(log.path) state['start_lr'] = state['lr'] print(state) imagenet_mean = [0.485, 0.456, 0.406] imagenet_std = [0.229, 0.224, 0.225] train_dataset = ImageList(args.root_folder, args.train_listfile, transform=transforms.Compose([ transforms.Resize(256), transforms.RandomCrop(224), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize(imagenet_mean, imagenet_std) ])) train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, pin_memory=False, num_workers=args.num_workers) val_dataset = ImageList(args.root_folder, args.val_listfile, transform=transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize(imagenet_mean, imagenet_std) ])) val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=args.batch_size, shuffle=False, pin_memory=False, num_workers=args.num_workers) if args.attention_depth == 0: from models.wide_resnet import WideResNet model = WideResNet().finetune(args.nlabels).cuda() else: from models.wide_resnet_attention import WideResNetAttention model = WideResNetAttention(args.nlabels, args.attention_depth, args.attention_width, args.has_gates, args.reg_weight).finetune(args.nlabels) # if args.load != "": # net.load_state_dict(torch.load(args.load), strict=False) # net = net.cuda() optimizer = optim.SGD([{ 'params': model.get_base_params(), 'lr': args.lr * 0.1 }, { 'params': model.get_classifier_params() }], lr=args.lr, weight_decay=1e-4, momentum=0.9, nesterov=True) if args.ngpu > 1: model = torch.nn.DataParallel(model, range(args.ngpu)).cuda() else: model = model.cuda() criterion = torch.nn.NLLLoss().cuda() def train(): """ """ model.train() for data, label in train_loader: data, label = torch.autograd.Variable(data, requires_grad=False).cuda(async=True), \ torch.autograd.Variable(label, requires_grad=False).cuda() optimizer.zero_grad() if args.attention_depth > 0: output, loss = model(data) if args.reg_weight > 0: loss = loss.mean() else: loss = 0 else: loss = 0 output = model(data) loss += F.nll_loss(output, label) loss.backward() optimizer.step() train_loss_meter.update(loss.data[0], data.size(0)) state['train_loss'] = train_loss_meter.mean() def val(): """ """ model.eval() for data, label in val_loader: data, label = torch.autograd.Variable(data, volatile=True).cuda(async=True), \ torch.autograd.Variable(label, volatile=True).cuda() if args.attention_depth > 0: output, loss = model(data) else: output = model(data) loss = F.nll_loss(output, label) val_loss_meter.update(loss.data[0], data.size(0)) preds = output.max(1)[1] val_accuracy_meter.update((preds == label).float().sum().data[0], data.size(0)) state['val_loss'] = val_loss_meter.mean() state['val_accuracy'] = val_accuracy_meter.mean() best_accuracy = 0 counter = 0 for epoch in range(args.epochs): train() val() harakiri.update(epoch, state['val_accuracy']) if state['val_accuracy'] > best_accuracy: counter = 0 best_accuracy = state['val_accuracy'] if args.save: torch.save(model.state_dict(), os.path.join(state["exp_dir"], "model.pytorch")) else: counter += 1 state['epoch'] = epoch + 1 log.update(state) print(state) if (epoch + 1) in args.schedule: for param_group in optimizer.param_groups: param_group['lr'] *= 0.1 state['lr'] *= 0.1