def __init__(self, args): super(GraphResNet, self).__init__() self.args = args # load or create encoder for initial features of nodes if args.pretrained_encoder_gnn is not None: self.encoder = ResNet18(output_layer=True) self.encoder.load_state_dict( torch.load(args.pretrained_encoder_gnn, map_location=args.device)) self.encoder.linear = Identity() # for params in self.encoder.parameters(): # params.requires_grad = False else: self.encoder = ResNet18(output_layer=False) embed_dim = 512 # output from ResNet18 # load pretrained contrastive model for computing similarity scores if args.pretrained_contrastive_model is not None: self.CoNet = ContrastiveNet(args) self.CoNet.load_state_dict( torch.load(args.pretrained_contrastive_model, map_location=args.device)) # for params in self.CoNet.parameters(): # params.requires_grad = False else: raise NotImplementedError( "only support pretrained contrastive_model") self.num_classes = 10 self.n = 10 # number of most similar nodes to consider for sending graph message self.graph_conv = WGraphConv(embed_dim, self.num_classes, bias=True)
def create_model(self): local_model = ResNet18(name='Local', created_time=self.params['current_time']) #device = torch.device("cpu") #local_model.to(device) local_model.cuda() target_model = ResNet18(name='Target', created_time=self.params['current_time']) #device = torch.device("cpu") #local_model.to(device) target_model.cuda() if self.params['resumed_model']: loaded_params = torch.load( f"saved_models/{self.params['resumed_model']}") target_model.load_state_dict(loaded_params['state_dict']) self.start_epoch = loaded_params['epoch'] self.params['lr'] = loaded_params.get('lr', self.params['lr']) logger.info( f"Loaded parameters from saved model: LR is" f" {self.params['lr']} and current epoch is {self.start_epoch}" ) else: self.start_epoch = 1 self.local_model = local_model self.target_model = target_model
def __init__(self, args): super(ContrastiveNet, self).__init__() # load or create encoder if args.pretrained_encoder_contrastive is not None: self.encoder = ResNet18(output_layer=True) self.encoder.load_state_dict( torch.load(args.pretrained_encoder_contrastive, map_location=args.device)) self.encoder.linear = Identity() else: self.encoder = ResNet18(output_layer=False) embed_dim = 512 # output from ResNet18 # We use a MLP with one hidden layer to obtain projection z_i = g(h_i) = W(2)σ(W(1)h_i) where σ is a ReLU non-linearity. self.projector = nn.Sequential( nn.Linear(embed_dim, embed_dim, bias=False), nn.ReLU(), nn.Linear(embed_dim, args.projection_dim, bias=False), ) self.device = args.device self.num_neg = args.num_neg self.batch_size = args.batch_size self.temperature = args.temperature self.device = args.device self.criterion = nn.CrossEntropyLoss(reduction="mean")
def initialize_model(model_name): ''' Initialise a model with a custom head to predict both sequence length and digits Parameters ---------- model_name : str Model Name can be either: ResNet VGG BaselineCNN ConvNet BaselineCNN_dropout Returns ------- model : object The model to be initialize ''' if model_name[:3] == "VGG": model = VGG(model_name, num_classes=7) model.classifier = CustomHead(512) elif model_name[:6] == "ResNet": if model_name == "ResNet18": model = ResNet18(num_classes=7) model.linear = CustomHead(512) elif model_name == "ResNet34": model = ResNet18(num_classes=7) model.linear = CustomHead(512) elif model_name == "ResNet50": model = ResNet50(num_classes=7) model.linear = CustomHead(512 * 4) elif model_name == "ResNet101": model = ResNet101(num_classes=7) model.linear = CustomHead(512 * 4) elif model_name == "ResNet152": model = ResNet152(num_classes=7) model.linear = CustomHead(512 * 4) elif model_name == "BaselineCNN": model = BaselineCNN(num_classes=7) model.fc2 = CustomHead(4096) elif model_name == "BaselineCNN_dropout": model = BaselineCNN_dropout(num_classes=7, p=0.5) model.fc2 = CustomHead(4096) return model
def __init__(self, base_net='resnet'): super(ResidualPhysics, self).__init__() if base_net == 'resnet': self.rgb_net = ResNet18(3, num_blocks=[2, 2, 2], symmetric=True) self.nir_net = ResNet18(1, num_blocks=[2, 2, 2], symmetric=True) # self.rgb_physics = nn.Conv2d(3, 256, kernel_size=3, stride=1, padding=1) # self.nir_physics = nn.Conv2d(1, 256, kernel_size=3, stride=1, padding=1) elif base_net == 'unet': pass else: raise RuntimeError('No such option!') self.frontbone = nn.Sequential( nn.Conv2d(4, 64, kernel_size=3, stride=1, padding=1), nn.ReLU(), nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1), nn.Conv2d(64, 3, kernel_size=3, stride=1, padding=1), nn.ReLU())
def __init__(self, model='VGG19', main_dir=main_dir_path, face_detector='undefined', use_cuda=False, reliability=0.8): self.main_dir = main_dir self.face_detector = face_detector self.use_cuda = use_cuda self.reliability = reliability self.cut_size = 44 self.transform_test = transforms.Compose([ transforms.TenCrop(self.cut_size), transforms.Lambda(lambda crops: torch.stack( [transforms.ToTensor()(crop) for crop in crops])), ]) self.class_names = [ 'Angry', 'Disgust', 'Fear', 'Happy', 'Sad', 'Surprise', 'Neutral' ] if model == 'VGG19': self.net = VGG('VGG19') elif model == 'Resnet18': self.net = ResNet18() self.checkpoint = torch.load(os.path.join( self.main_dir + 'pretrained_model/' + model, 'PrivateTest_model.t7'), map_location='cpu') self.net.load_state_dict(self.checkpoint['net']) if self.use_cuda: self.net.cuda() self.net.eval()
def run_stochastic_Facloc(data, targets, budget): model = ResNet18(num_cls) model = model.to(device) approximate_error = 0.01 per_iter_bud = 10 num_iterations = int(budget / 10) facloc_indices = [] trn_indices = list(np.arange(len(data))) sample_size = int( len(data) / num_iterations * math.log(1 / approximate_error)) #greedy_batch_size = 1200 for i in range(num_iterations): print(i) rem_indices = list(set(trn_indices).difference(set(facloc_indices))) state = np.random.get_state() np.random.seed(i * i) sub_indices = np.random.choice(rem_indices, size=sample_size, replace=False) np.random.set_state(state) data_subset = data[sub_indices].cpu() #targets_subset = targets[sub_indices].cpu() #train_loader_greedy = [] #train_loader_greedy.append((data_subset, targets_subset)) setf_model = SetFunctionFacLoc(device, 800) idxs = setf_model.lazy_greedy_max(per_iter_bud, data_subset, model) facloc_indices.extend([sub_indices[idx] for idx in idxs]) return facloc_indices
def train_cifar_models(): # ResNet training net = ResNet18() net.cuda() train_cifar10_model(net, learning_rates=[0.001, 0.0001], iters=[50, 30], output_path='models/resnet18_cifar10.model') # Cifar Tiny net = Cifar_Tiny() net.cuda() train_cifar10_model(net, learning_rates=[0.001, 0.0001], iters=[50, 30], output_path='models/tiny_cifar10.model') # Cifar Very Tiny net = Cifar_Very_Tiny() net.cuda() train_cifar10_model(net, learning_rates=[ 0.001, ], iters=[ 50, ], output_path='models/very_tiny_cifar10.model')
def get_model(model_name, num_classes=512): if model_name == 'VGG19': model = VGG('VGG19') elif model_name == 'ResNet18': model = ResNet18(num_classes) elif model_name == 'PreActResNet18': model = PreActResNet18() elif model_name == 'GoogLeNet': model = GoogLeNet() elif model_name == 'DenseNet121': model = DenseNet121() elif model_name == 'ResNeXt29_2x64d': model = ResNeXt29_2x64d() elif model_name == 'MobileNet': model = MobileNet() elif model_name == 'MobileNetV2': model = MobileNetV2() elif model_name == 'DPN92': model = DPN92() elif model_name == 'ShuffleNetG2': model = ShuffleNetG2() elif model_name == 'SENet18': model = SENet18() elif model_name == 'ShuffleNetV2': model = ShuffleNetV2(1) elif model_name == 'EfficientNetB0': model = EfficientNetB0() else: print('{} does NOT exist in repertory.'.format(model_name)) sys.exit(1) return model
def main(): global args args = parser.parse_args() print(args) # create model print("=> creating model '{}'".format(args.arch)) if args.arch == 'vgg16': from models.vgg import VGG model = nn.DataParallel(VGG('VGG16', nclass=args.num_classes), device_ids=range(1)) if args.arch == 'resnet18': from models.resnet import ResNet18 model = nn.DataParallel(ResNet18().cuda()) checkpoint = torch.load('./checkpoint/cifar10_vgg16_teacher.pth') model.load_state_dict(checkpoint) cudnn.benchmark = True # Data loading code mix_img_val_loader = torch.utils.data.DataLoader( active_query_loader(transforms.Compose([ transforms.ToTensor(), ])), batch_size=args.batch_size, shuffle=False, num_workers=args.workers, pin_memory=True) ## image extraction logits_extract(mix_img_val_loader, model)
def get_model(arch, num_classes, channels=3): """ Args: arch: string, Network architecture num_classes: int, Number of classes channels: int, Number of input channels Returns: model, nn.Module, generated model """ if arch.lower() == "resnet18": model = ResNet18(channels, num_classes) elif arch.lower() == "resnet34": model = ResNet34(channels, num_classes) elif arch.lower() == "resnet50": model = ResNet50(channels, num_classes) elif arch.lower() == "resnet101": model = ResNet101(channels, num_classes) elif arch.lower() == "resnet152": model = ResNet152(channels, num_classes) elif arch.lower() == "mobilenet_v1": model = MobileNetV1(num_classes, channels) elif arch.lower() == "mobilenet_v2": model = MobileNetV2(num_classes, channels) else: raise NotImplementedError( f"{arch} not implemented. " f"For supported architectures see documentation") return model
def evaluate_cifar_models_retrieval(): """ Evaluates the baselines teacher/students :return: """ evaluate_model_retrieval( net=Cifar_Tiny(num_classes=10), path='models/tiny_cifar10.model', result_path='results/tiny_cifar10_baseline_retrieval.pickle', layer=3) evaluate_model_retrieval( net=Cifar_Tiny(num_classes=10), path='models/tiny_cifar10.model', result_path='results/tiny_cifar10_baseline_retrieval_e.pickle', layer=3, metric='l2') evaluate_model_retrieval( net=ResNet18(num_classes=10), path='models/resnet18_cifar10.model', result_path='results/resnet18_cifar10_baseline_retrieval.pickle', layer=3) evaluate_model_retrieval( net=ResNet18(num_classes=10), path='models/resnet18_cifar10.model', result_path='results/resnet18_cifar10_baseline_retrieval_e.pickle', layer=3, metric='l2') evaluate_model_retrieval( net=Cifar_Very_Tiny(num_classes=10), path='models/very_tiny_cifar10.model', result_path='results/very_tiny_cifar10_baseline_retrieval.pickle', layer=3) evaluate_model_retrieval( net=Cifar_Very_Tiny(num_classes=10), path='models/very_tiny_cifar10.model', result_path='results/very_tiny_cifar10_baseline_retrieval_e.pickle', layer=3, metric='l2')
def load_mod_dl(args): """ :param args: :return: """ if args.dataset == 'cifar10': imsize, in_channel, num_classes = 32, 3, 10 train_loader, val_loader, test_loader = data_loaders.load_cifar10( args.batch_size, val_split=True, augmentation=args.data_augmentation, subset=[args.train_size, args.val_size, args.test_size]) elif args.dataset == 'cifar100': imsize, in_channel, num_classes = 32, 3, 100 train_loader, val_loader, test_loader = data_loaders.load_cifar100( args.batch_size, val_split=True, augmentation=args.data_augmentation, subset=[args.train_size, args.val_size, args.test_size]) elif args.dataset == 'mnist': imsize, in_channel, num_classes = 28, 1, 10 num_train = 50000 train_loader, val_loader, test_loader = data_loaders.load_mnist( args.batch_size, subset=[args.train_size, args.val_size, args.test_size], num_train=num_train, only_split_train=False) if args.model == 'resnet18': cnn = ResNet18(num_classes=num_classes) elif args.model == 'cbr': cnn = CBRStudent(in_channel, num_classes) # This essentially does no mixup. mixup_mat = -100 * torch.ones([num_classes, num_classes]).cuda() checkpoint = None if args.load_checkpoint: checkpoint = torch.load(args.load_checkpoint) mixup_mat = checkpoint['mixup_grid'] print(f"loaded mixupmat from {args.load_checkpoint}") if args.rand_mixup: # Randomise mixup grid rng = np.random.RandomState(args.seed) mixup_mat = rng.uniform( 0.5, 1.0, (num_classes, num_classes)).astype(np.float32) print("Randomised the mixup mat") mixup_mat = torch.from_numpy( mixup_mat.reshape(num_classes, num_classes)).cuda() model = cnn.cuda() model.train() return model, mixup_mat, train_loader, val_loader, test_loader, checkpoint
def get_model_for_training(model_name, device): if model_name == 'ResNet18': print("Model: ResNet18", file=sys.stderr) return ResNet18().to(device), ResNet18().to(device) elif model_name == 'ResNet34': print("Model: ResNet34", file=sys.stderr) return ResNet34().to(device), ResNet34().to(device) elif model_name == 'PreActResNet18': print("Model: PreActivate ResNet18", file=sys.stderr) return PreActResNet18().to(device), PreActResNet18().to(device) elif model_name == 'PreActResNet34': print("Model: PreActivate ResNet 34") return PreActResNet34().to(device), PreActResNet34().to(device) elif model_name == 'WideResNet28': print("Model: Wide ResNet28", file=sys.stderr) return WideResNet28().to(device), WideResNet28().to(device) elif model_name == 'WideResNet34': print("Model: Wide ResNet34", file=sys.stderr) return WideResNet34().to(device), WideResNet34().to(device)
def build_model(cf): def quantized_relu(x): return quantize_op(x,nb=cf.abits) H = 1. if cf.network_type =='float': Conv = Conv2D Fc = Dense Act = lambda: LeakyReLU() elif cf.network_type in ['qnn', 'full-qnn']: Conv = lambda **kwargs: QuantizedConv2D(H=1, nb=cf.wbits, **kwargs) Fc = lambda **kwargs: QuantizedDense(nb=cf.abits, **kwargs) if cf.network_type=='qnn': Act = lambda: LeakyReLU() else: # full-qnn Act = lambda: Activation(quantized_relu) elif cf.network_type in ['bnn', 'qbnn', 'full-bnn']: Conv = lambda **kwargs: BinaryConv2D(H=1, **kwargs) Fc = BinaryDense if cf.network_type=='bnn': Act = lambda: LeakyReLU() elif cf.network_type=='qbnn': Act = lambda: Activation(quantized_relu) else: #full-bnn Act = lambda: Activation(binary_tanh) elif cf.network_type in ['tnn', 'qtnn', 'full-tnn']: Conv = lambda **kwargs: TernaryConv2D(H=1, **kwargs) Fc = TernaryDense if cf.network_type=='tnn': Act = lambda: LeakyReLU() elif cf.network_type=='qtnn': Act = lambda: Activation(quantized_relu) else: #full-tnn Act = lambda: Activation(ternary_tanh) else: raise ValueError('wrong network type, the supported network types in this repo are float, qnn, full-qnn, bnn and full-bnn') if cf.architecture=="VGG": model = Vgg(Conv, Act, Fc, cf) elif cf.architecture=="RESNET": model = ResNet18(Conv, Act, Fc, cf) else: raise ValueError("Error: type " + str(cf.architecture) + " is not supported") model.summary() return model
def get_new_model(tmp_scale = True, num_classes = args.num_classes): if args.model == 'resnet18': return ResNet18(tmp_scale = tmp_scale, num_classes = num_classes) elif args.model == 'resnet50': return ResNet50(tmp_scale = tmp_scale, num_classes = num_classes) elif args.model == 'resnet101': return ResNet101(tmp_scale = tmp_scale, num_classes = num_classes) elif args.model == 'inceptionv4': return inceptionv4(tmp_scale = tmp_scale, num_classes = num_classes) elif args.model == 'densenet': return DenseNet(tmp_scale = tmp_scale)
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_model(args): if args.model == "ResNet": from models.resnet import ResNet18 model = ResNet18(output_layer=True, num_classes=10) elif args.model == "GraphResNet": from models.model import GraphResNet model = GraphResNet(args) else: raise NotImplementedError model = model.to(args.device) return model
def train(args): base_dir = args.base_dir num_classes = args.num_classes model_name = str.lower(args.model) if model_name== 'lenet_baseline': input_shape = (32, 32, 3) model = LeNet_baseline(input_shape=input_shape, num_classes=num_classes) elif model_name == 'lenet_modified': input_shape = (32, 32, 3) model = LeNet_modified(input_shape=input_shape, num_classes=num_classes) elif model_name == 'alexnet': input_shape = (227, 227, 3) model = AlexNet(input_shape=input_shape, num_classes=num_classes) elif model_name == 'vgg16': input_shape = (224, 224, 3) model = VGG16(input_shape=input_shape, num_classes=num_classes) elif model_name == 'vgg19': input_shape = (224, 224, 3) model = VGG19(input_shape=input_shape, num_classes=num_classes) elif model_name == 'resnet18': input_shape = (112, 112, 3) model = ResNet18(input_shape=input_shape, num_classes=num_classes) else: print('Please choose an implemented model!') sys.exit() # get training dataset x, y = input_preprocess(base_dir, input_shape, num_classes) # split dataset into train and val set x_train, x_val, y_train, y_val = train_test_split(x, y, test_size=0.2, random_state=seed) print('Train images loaded') # one-hot encoding for each label y_train = one_hot(y_train, num_classes) y_val = one_hot(y_val, num_classes) # train the model callbacks = get_callbacks(args.checkpoint, model_name) model.compile(optimizer=Adam(lr=args.lr), loss='categorical_crossentropy', metrics=['acc']) model.summary() history = model.fit(x_train, y_train, batch_size=args.batch_size, validation_data=(x_val, y_val), epochs=args.epochs, callbacks=callbacks) visualize(history) return model
def create_model(self): local_model = ResNet18(name='Local', created_time=self.params['current_time']) local_model.to(self.device) target_model = ResNet18(name='Target', created_time=self.params['current_time']) target_model.to(self.device) if self.resumed_model: loaded_params = torch.load( f"{self.params['repo_path']}/saved_models/{self.params['resumed_model']}" ) target_model.load_state_dict(loaded_params['state_dict']) self.start_round = loaded_params['round'] self.params['lr'] = loaded_params.get('lr', self.params['lr']) logger.info( f"Loaded parameters from saved model: LR is" f" {self.params['lr']} and current round is {self.start_round}" ) else: self.start_round = 1 self.local_model = local_model self.target_model = target_model
def get_model(model='ResNet18', **kwargs): if model == 'ResNet18': return ResNet18(n_classes=kwargs.get('n_classes', 10)) elif 'VGG' in model: return VGG( vgg_name=model, k=kwargs['k'], dropout=kwargs.get('dropout', None), n_classes=kwargs.get('n_classes', 10), ) elif 'LeNet' == model: return LeNet(dropout=kwargs.get('dropout', None)) else: raise NotImplementedError('unknown {} model'.format(model))
def get_classifier(mode, n_classes=10): if mode == 'resnet18': classifier = ResNet18(num_classes=n_classes) elif mode == 'resnet34': classifier = ResNet34(num_classes=n_classes) elif mode == 'resnet50': classifier = ResNet50(num_classes=n_classes) elif mode == 'resnet18_imagenet': classifier = resnet18(num_classes=n_classes) elif mode == 'resnet50_imagenet': classifier = resnet50(num_classes=n_classes) else: raise NotImplementedError() return classifier
def load_model_bold(args): if args.model == "ResNet": from models.resnet import ResNet18 model = ResNet18(output_layer=True, num_classes=10) elif args.model == "ResNetGCN": from models.model import ResNetGCN model = ResNetGCN(args.graph_type) elif args.model == "ResNetGATCN": from models.model import ResNetGATCN model = ResNetGATCN(args.graph_type) else: raise NotImplementedError model = model.to(args.device) return model
def run_transfer(learning_rates=(0.001, 0.0001), iters=(3, 0), method='mds'): torch.manual_seed(12345) student_layers, teacher_layers, weights, loss_params, T = (3,), (3,), (1,), {}, 2 print(method) transfer_name = method # Output paths output_path = 'models/aux_' + transfer_name + '.model' results_path = 'results/aux_' + transfer_name # Load a pre-trained teacher network student_net = Cifar_Tiny(10) # Use a pre-trained model load_model(student_net, 'models/tiny_cifar10.model') # Load the teacher model teacher_net = ResNet18(num_classes=10) load_model(teacher_net, 'models/resnet18_cifar10.model') train_loader, test_loader, train_loader_raw = cifar10_loader(batch_size=128) # Move the models into GPU student_net.cuda() teacher_net.cuda() # Perform the transfer W = None for lr, iters in zip(learning_rates, iters): if method == 'pkt': kernel_parameters = {'student': 'combined', 'teacher': 'combined', 'loss': 'combined'} prob_transfer(student_net, teacher_net, train_loader, epochs=iters, lr=lr, teacher_layers=teacher_layers, student_layers=student_layers, layer_weights=weights, kernel_parameters=kernel_parameters, loss_params=loss_params) else: assert False save_model(student_net, output_path) print("Model saved at ", output_path) # Perform the evaluation evaluate_model_retrieval(net=Cifar_Tiny(num_classes=10), path=output_path, result_path=results_path + '_retrieval.pickle', layer=3) evaluate_model_retrieval(net=Cifar_Tiny(num_classes=10), path=output_path, result_path=results_path + '_retrieval_e.pickle', layer=3, metric='l2')
def get_classifier(mode, n_classes=10): if mode == 'resnet18': classifier = ResNet18(num_classes=n_classes) elif mode == 'resnet34': classifier = ResNet34(num_classes=n_classes) elif mode == 'resnet50': classifier = ResNet50(num_classes=n_classes) elif mode == 'resnet18_imagenet': classifier = resnet18(num_classes=n_classes) elif mode == 'resnet50_imagenet': classifier = resnet50(num_classes=n_classes) elif mode == 'live': classifier = FeatherNet(input_size=128, se=True, avgdown=True) # classifier = ResNet18(num_classes=n_classes) # classifier = LiveModel() else: raise NotImplementedError() return classifier
def load_mod_dl(args): """ :param args: :return: """ if args.dataset == 'cifar10': imsize, in_channel, num_classes = 32, 3, 10 train_loader, val_loader, test_loader = data_loaders.load_cifar10(args.batch_size, val_split=True, augmentation=args.data_augmentation, subset=[args.train_size, args.val_size, args.test_size]) elif args.dataset == 'cifar100': imsize, in_channel, num_classes = 32, 3, 100 train_loader, val_loader, test_loader = data_loaders.load_cifar100(args.batch_size, val_split=True, augmentation=args.data_augmentation, subset=[args.train_size, args.val_size, args.test_size]) elif args.dataset == 'mnist': imsize, in_channel, num_classes = 28, 1, 10 num_train = 50000 train_loader, val_loader, test_loader = data_loaders.load_mnist(args.batch_size, subset=[args.train_size, args.val_size, args.test_size], num_train=num_train, only_split_train=False) if args.model == 'resnet18': cnn = ResNet18(num_classes=num_classes, num_channels=in_channel) elif args.model == 'cbr': cnn = CBRStudent(in_channel, num_classes) mixup_mat = -1*torch.ones([num_classes,num_classes]).cuda() mixup_mat.requires_grad = True 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, mixup_mat, train_loader, val_loader, test_loader, checkpoint
def get_model(model): model_path = '../saved' if model == 'LeNet-5': net = LeNet() model_name = 'lenet.pth' elif model == 'VGG-16': net = Vgg16_Net() model_name = 'vgg16.pth' elif model == 'ResNet18': net = ResNet18() model_name = 'resnet18.pth' elif model == 'ResNet34': net = ResNet34() model_name = 'resnet34.pth' elif model == 'ResNet50': net = ResNet50() model_name = 'resnet50.pth' else: net = ResNet101() model_name = 'resnet101.pth' return net, os.path.join(model_path, model_name)
def get_model(args): if args.dataset == 'cifar-10': num_classes = 10 elif args.dataset == 'cifar-100': num_classes = 100 else: raise NotImplementedError if 'contrastive' in args.train_type or 'linear_eval' in args.train_type: contrastive_learning = True else: contrastive_learning = False if args.model == 'ResNet18': model = ResNet18(num_classes, contrastive_learning) print('ResNet18 is loading ...') elif args.model == 'ResNet50': model = ResNet50(num_classes, contrastive_learning) print('ResNet 50 is loading ...') return model
def train_cifar_models(): """ Trains the baselines teacher/students :return: """ # ResNet training net = ResNet18() net.cuda() train_cifar10_model(net, learning_rates=[0.001, 0.0001], iters=[50, 50], output_path='models/resnet18_cifar10.model') # Cifar Tiny net = Cifar_Tiny() net.cuda() train_cifar10_model(net, learning_rates=[0.001, 0.0001], iters=[50, 50], output_path='models/tiny_cifar10.model')
def model_init(self, args): # Network if args.dataset == 'MNIST': print("MNIST") self.net = cuda(ToyNet_MNIST(y_dim=self.y_dim), self.cuda) elif args.dataset == 'CIFAR10': print("Dataset used CIFAR10") if args.network_choice == 'ToyNet': self.net = cuda(ToyNet_CIFAR10(y_dim=self.y_dim), self.cuda) elif args.network_choice == 'ResNet18': self.net = cuda(ResNet18(), self.cuda) elif args.network_choice == 'ResNet34': self.net = cuda(ResNet34(), self.cuda) elif args.network_choice == 'ResNet50': self.net = cuda(ResNet50(), self.cuda) self.net.weight_init(_type='kaiming') # setup optimizer self.optim = optim.Adam([{ 'params': self.net.parameters(), 'lr': self.lr }], betas=(0.5, 0.999))