hi, i am facing this problem and i don’t know how to solve it.
import pickle
import torch
import torch.nn as nn
import torch.optim as optim
from pandas.plotting._matplotlib import hist
from torchvision.transforms import transforms
from torch.utils.data import DataLoader
import torchvision.datasets as datasets
import numpy as np
import pandas as pd
import sys
import json
import os
import sklearn.metrics as metrics
from watermark_regularizer import WatermarkRegularizer
from watermark_regularizer import get_wmark_regularizer
from resnet import ResNet
set device
device = torch.device(‘cuda’ if torch.cuda.is_available() else ‘cpu’)
RESULT_PATH = ‘./result’
MODEL_CHKPOINT_FNAME = os.path.join(RESULT_PATH, ‘WRN-Weights.pth’)
def update_hdf5(fname, path, data):
store = pd.HDFStore(fname)
if path in store.keys():
store.remove(path)
store.append(path, data)
store.close()
def save_wmark_signatures(prefix, model):
for layer_id, wmark_regularizer in get_wmark_regularizer(model):
fname_w = prefix + ‘_layer{}_w.npy’.format(layer_id)
fname_b = prefix + ‘_layer{}_b.npy’.format(layer_id)
np.save(fname_w, wmark_regularizer.get_matrix())
np.save(fname_b, wmark_regularizer.get_signature())
lr_schedule = [60, 120, 160]
def schedule(epoch_idx):
if(epoch_idx + 1) < lr_schedule[0]:
return 0.1
elif(epoch_idx + 1) < lr_schedule[1]:
return 0.02
elif(epoch_idx + 1) < lr_schedule[2]:
return 0.004
return 0.0008
if name == ‘main’:
if len(sys.argv) < 2:
print(“Fatal: You forgot to include the directory name on the command line.”)
print(“Usage: python %s ” % sys.argv[0])
sys.exit(1)
settings_json_fname = sys.argv[1]
train_settings = json.load(open(settings_json_fname))
if not os.path.isdir(RESULT_PATH):
os.makedirs(RESULT_PATH)
torch.autograd.set_detect_anomaly(True) # enable anomaly detection
# load dataset and fitting data for learning
if train_settings['dataset'] == 'cifar10':
transform_train = transforms.Compose([transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
])
transform_test = transforms.Compose([transforms.ToTensor(),
transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
])
train_set = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform_train)
test_set = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform_test)
train_loader = DataLoader(train_set, batch_size=train_settings['batch_size'], shuffle=True)
test_loader = DataLoader(test_set, batch_size=train_settings['batch_size'], shuffle=False)
nb_classes = 10
else:
print('not supported dataset"{}"'.format(train_settings['dataset']))
exit(1)
if 'replace_train_y' in train_settings and len(train_settings['replace_train_y']) > 0:
print('train_loader was replaced from"{}"'.format(train_settings['replace_train_y']))
# train_loader = np.load(train_settings['replace_train_y'])
train_loader = torch.from_numpy(np.load(train_settings['replace_train_y']))
# read parameters
batch_size = train_settings['batch_size']
nb_epoch = train_settings['epoch']
scale = train_settings['scale']
embed_dim = train_settings['embed_dim']
N = train_settings['N']
k = train_settings['k']
target_blk_id = train_settings['target_blk_id']
base_modelw_fname = train_settings['base_modelw_fname']
wtype = train_settings['wmark_wtype']
randseed = train_settings['randseed'] if 'randseed' in train_settings else 'none'
ohist_fname = train_settings['history']
hist_hdf_path = 'WTYPE_{}/DIM{}/SCALE{}/N{}K{}B{}EPOCH{}/TBLK{}'.format(wtype, embed_dim, scale, N, k, batch_size, nb_epoch, target_blk_id)
modelname_prefix = os.path.join(RESULT_PATH, 'wrn_' + hist_hdf_path.replace('/', '_'))
# initialize process for watermark
b = np.ones((1, embed_dim))
wmark_regularizer = WatermarkRegularizer(scale, b, wtype=wtype, randseed=randseed)
init_shape = (3, 32, 32)
# model = ResNet(init_shape, nb_classes)
# model = ResNet(init_shape, image_channels=3, num_classes=10, layers=3)
model = ResNet(image_channels=3, num_classes=10, layers=[2, 3, 4])
print(model)
print('watermark matrix: \n{}'.format(wmark_regularizer.get_matrix()))
# training process
optimizer = optim.SGD(model.parameters(), lr=0.1, momentum=0.9, nesterov=True)
criterion = nn.CrossEntropyLoss()
scheduler = optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda x: schedule(x))
if len(base_modelw_fname) > 0:
model.load_state_dict(torch.load(base_modelw_fname))
print("finished compiling")
# train network
for epoch in range(nb_epoch):
losses = []
for batch_idx, (data, targets) in enumerate(train_loader):
# get data to cuda if possible
data = data.to(device=device)
targets = targets.to(device=device)
# forward
scores = model(data)
loss = criterion(scores, targets)
losses.append(loss.item())
# backward
optimizer.zero_grad()
loss.backward()
# gradient descent or adam step
optimizer.step()
print(f'Cost at epoch {epoch} is {sum(losses) / len(losses)}')
# check accuracy
def check_accuracy(loader, model):
if loader.dataset:
print("checking accuracy on training data")
else:
print("checking accuracy on test data")
num_correct = 0
num_samples = 0
model.eval()
with torch.no_grad():
for x, y in loader:
x = x.to(device=device)
y = y.to(device=device)
scores = model(x)
_, predictions = scores.max(1)
num_correct += (predictions == y).sum()
num_samples += predictions.size(0)
print(f'got {num_correct}/{num_samples} with accuracy {float(num_correct)/float(num_samples)*100:2f}')
model.train()
check_accuracy(train_loader, model)
check_accuracy(test_loader, model)
####
torch.save(model.state_dict(), modelname_prefix + '.pt')
with open(ohist_fname, 'ab') as f:
pickle.dump(hist.history, f)
if target_blk_id > 0:
save_wmark_signatures(model, modelname_prefix)