import caffe
from caffe import layers as L, params as P

# 设置计算模式
caffe.set_mode_gpu()  # 或 caffe.set_mode_cpu()

# 加载预训练模型和权重
model_def = 'deploy.prototxt'  # 模型定义文件
model_weights = 'pretrained.caffemodel'  # 预训练权重

net = caffe.Net(model_def, model_weights, caffe.TRAIN)

# 示例：修改分类层以适应新的类别数
n_new_classes = 10  # 新的类别数

# 获取最后一层的名称
last_layer = list(net._layer_names)[-1]

# 修改全连接层的输出维度
net.params[last_layer][0].data[...] = np.random.randn(n_new_classes, net.params[last_layer][0].data.shape[1]) * 0.01
net.params[last_layer][1].data[...] = np.zeros(n_new_classes)

# 创建数据转换器
transformer = caffe.io.Transformer({'data': net.blobs['data'].data.shape})
transformer.set_transpose('data', (2,0,1))  # 从HWC转为CHW
transformer.set_mean('data', np.array([104, 117, 123]))  # 减去均值
transformer.set_raw_scale('data', 255)  # 从[0,1]转为[0,255]
transformer.set_channel_swap('data', (2,1,0))  # 从RGB转为BGR

# 创建求解器
solver = caffe.SGDSolver('solver.prototxt')

# 复制预训练权重
solver.net.copy_from(model_weights)

# 设置基础学习率
base_lr = 0.001
solver.param.base_lr = base_lr

# 设置不同层的学习率
for layer_name in solver.net.params:
    if 'conv' in layer_name:
        solver.net.params[layer_name][0].lr = base_lr * 1
        solver.net.params[layer_name][1].lr = base_lr * 2
    elif 'fc' in layer_name:
        solver.net.params[layer_name][0].lr = base_lr * 10
        solver.net.params[layer_name][1].lr = base_lr * 20

# 训练参数
max_iter = 10000
display = 100

for it in range(max_iter):
    solver.step(1)  # 一次迭代
    
    if it % display == 0:
        # 计算并显示损失和准确率
        loss = solver.net.blobs['loss'].data
        acc = solver.net.blobs['accuracy'].data
        print(f'Iteration {it}, loss={loss}, accuracy={acc}')