用pytorch写一个域适应迁移学习代码，损失函数为mmd距离域判别损失和交叉熵损失

可以使用以下代码实现：

import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
from sklearn.metrics import accuracy_score
from torch.autograd import Variable

class DomainAdaptationModel(nn.Module):
    def __init__(self, num_classes=2):
        super(DomainAdaptationModel, self).__init__()
        self.feature_extractor = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=5),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(64, 128, kernel_size=5),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(128, 256, kernel_size=5),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(256, 512, kernel_size=5),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(512, 1024, kernel_size=5),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),
        )
        self.classifier = nn.Sequential(
            nn.Linear(1024, 1024),
            nn.ReLU(inplace=True),
            nn.Dropout(),
            nn.Linear(1024, num_classes),
        )

    def forward(self, x):
        x = self.feature_extractor(x)
        x = x.view(x.size(0), -1)
        x = self.classifier(x)
        return x

def mmd_loss(source_features, target_features):
    source_mean = torch.mean(source_features, dim=0)
    target_mean = torch.mean(target_features, dim=0)
    mmd = torch.mean(torch.pow(source_mean - target_mean, 2))
    return mmd

def domain_discriminator_loss(source_features, target_features):
    source_labels = torch.zeros(source_features.size(0))
    target_labels = torch.ones(target_features.size(0))
    labels = torch.cat((source_labels, target_labels), dim=0)
    features = torch.cat((source_features, target_features), dim=0)
    criterion = nn.BCEWithLogitsLoss()
    loss = criterion(features, labels)
    return loss

def train(model, source_loader, target_loader, optimizer, num_epochs=10):
    model.train()
    for epoch in range(num_epochs):
        for i, (source_data, target_data) in enumerate(zip(source_loader, target_loader)):
            source_inputs, source_labels = source_data
            target_inputs, _ = target_data
            inputs = torch.cat((source_inputs, target_inputs), dim=0)
            inputs = Variable(inputs.cuda())
            source_labels = Variable(source_labels.cuda())
            optimizer.zero_grad()
            source_features = model(inputs[:source_inputs.size(0)])
            target_features = model(inputs[source_inputs.size(0):])
            mmd_loss_value = mmd_loss(source_features, target_features)
            domain_discriminator_loss_value = domain_discriminator_loss(source_features, target_features)
            classification_loss = nn.CrossEntropyLoss()(model(inputs[:source_inputs.size(0)]), source_labels)
            loss = classification_loss + mmd_loss_value + domain_discriminator_loss_value
            loss.backward()
            optimizer.step()
            if i % 100 == 0:
                print('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}'.format(epoch+1, num_epochs, i+1, len(source_loader), loss.item()))

source_loader = torch.utils.data.DataLoader(source_dataset, batch_size=32, shuffle=True)
target_loader = torch.utils.data.DataLoader(target_dataset, batch_size=32, shuffle=True)
model = DomainAdaptationModel(num_classes=10)
model.cuda()
optimizer = optim.Adam(model.parameters(), lr=0.001)
train(model, source_loader, target_loader, optimizer, num_epochs=10)