def main():
# Set the hyper-parameters
learning_rate = 0.001
batch_size = 100
epoches = 11
train_loader, test_loader = MNIST_dataset(batch_size)
# Build the model
lenet = Lenet()
# Set the loss function & optimizer
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(lenet.parameters(), lr=learning_rate)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# Train & Test the model
training_loss = train_and_test(device,
epoches,
train_loader,
test_loader,
optimizer,
criterion,
model=lenet)
plot_loss(training_loss)
def active_learner(device, seed):
torch.backends.cudnn.deterministic = True
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
np.random.seed(seed)
random.seed(seed)
model = Lenet().to(device)
optimizer = torch.optim.Adam(model.parameters(),
lr=0.001,
amsgrad=True)
return model, optimizer, None
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,
exp_lr_scheduler,
num_epochs=100)
