def human_cnn(self): human_cnn = Lenet() human_cnn.load_state_dict( torch.load( 'utils/cifar10_2class/trained_models/humancnn/lenet_weights_0.5_normalisation', map_location=self.device)) human_cnn.to(self.device) return human_cnn
dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'val']} class_names = image_datasets['train'].classes print("class_names", class_names) model_ft = Lenet() model_ft.conv1 = nn.Conv2d(3, 6, 3) # 这里是彩色图像 num_ftrs = model_ft.fc3.in_features # Alternatively, it can be generalized to nn.Linear(num_ftrs, len(class_names)). model_ft.fc = nn.Linear(num_ftrs, len(class_names)) if torch.cuda.device_count() > 1: print("Let's use", torch.cuda.device_count(), "GPUs!") # dim = 0 [30, xxx] -> [10, ...], [10, ...], [10, ...] on 3 GPUs model_ft = nn.DataParallel(model_ft) model_ft = model_ft.to(device) criterion = nn.CrossEntropyLoss() # Observe that all parameters are being optimized # optimizer_ft = optim.SGD(model_ft.parameters(), lr=0.001, momentum=0.9) optimizer_ft = optim.Adam(model_ft.parameters()) # Decay LR by a factor of 0.1 every 7 epochs exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=7, gamma=0.1) model_ft = train_model(model_ft, criterion, optimizer_ft,