import torch
import torchvision
import torch.nn as nn
import torch.optim as optim
import numpy as np
from sklearn.metrics import roc_auc_score
# set up the device
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# define the data transformations
data_transforms = torchvision.transforms.Compose([
torchvision.transforms.Resize((224, 224)),
torchvision.transforms.ToTensor(),
torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
# set up the hyperparameters
lr = 1e-4
batch_size = 32
num_workers = 4
num_epochs = 10
epsilon = 0.1
# load source dataset
source_dataset = torchvision.datasets.ImageFolder(
root='数据集1路径',
transform=data_transforms
)
source_loader = torch.utils.data.DataLoader(
source_dataset,
batch_size=batch_size,
shuffle=True,
num_workers=num_workers
)
# load target dataset
target_dataset = torchvision.datasets.ImageFolder(
root='数据集2路径',
transform=data_transforms
)
target_loader = torch.utils.data.DataLoader(
target_dataset,
batch_size=batch_size,
shuffle=True,
num_workers=num_workers
)
# define the network
class ResNet18(nn.Module):
def __init__(self):
super(ResNet18, self).__init__()
self.resnet = torchvision.models.resnet18(pretrained=True)
self.resnet.fc = nn.Linear(512, 2)
def forward(self, x):
x = self.resnet(x)
return x
# initialize the network and the optimizer
model = ResNet18().to(device)
optimizer = optim.Adam(model.parameters(), lr=lr)
# define the loss function
criterion = nn.CrossEntropyLoss()
# train the model
for epoch in range(num_epochs):
running_loss = 0.0
for i, data in enumerate(source_loader, 0):
inputs, labels = data
inputs, labels = inputs.to(device), labels.to(device)
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
running_loss += loss.item()
# Calculate AUC for epoch
target_data = next(iter(target_loader))[0].to(device)
with torch.no_grad():
source_preds = model(inputs).detach().cpu().numpy()
target_preds = model(target_data).detach().cpu().numpy()
y_true = np.concatenate((np.zeros(len(source_preds)), np.ones(len(target_preds))))
y_score = np.concatenate((source_preds[:, 1], target_preds[:, 1]))
auc = roc_auc_score(y_true, y_score)
print(f"Epoch {epoch + 1}, AUC: {auc:.4f}")
print("Training is done!")
训练结果:AUC越高(一般认为大于0.7),模型识别能力越强,说明两个数据集之前的域分布差异较大;反之较小。