def __init__(self, T: int, channels: int, use_cupy=False):
super().__init__()
self.T = T
self.conv_fc = nn.Sequential(
layer.Conv2d(1, channels, kernel_size=3, padding=1, bias=False),
layer.BatchNorm2d(channels),
neuron.IFNode(surrogate_function=surrogate.ATan()),
layer.MaxPool2d(2, 2), # 14 * 14
layer.Conv2d(channels,
channels,
kernel_size=3,
padding=1,
bias=False),
layer.BatchNorm2d(channels),
neuron.IFNode(surrogate_function=surrogate.ATan()),
layer.MaxPool2d(2, 2), # 7 * 7
layer.Flatten(),
layer.Linear(channels * 7 * 7, channels * 4 * 4, bias=False),
neuron.IFNode(surrogate_function=surrogate.ATan()),
layer.Linear(channels * 4 * 4, 10, bias=False),
neuron.IFNode(surrogate_function=surrogate.ATan()),
)
functional.set_step_mode(self, step_mode='m')
if use_cupy:
functional.set_backend(self, backend='cupy')
def __init__(self):
super().__init__()
self.fc = nn.Sequential(
layer.Linear(28, 32),
neuron.IFNode(surrogate_function=surrogate.ATan()),
layer.Linear(32, 10),
neuron.IFNode(surrogate_function=surrogate.ATan())
)
def __init__(self):
super().__init__()
self.fc = nn.Sequential(
layer.Linear(28, 32),
neuron.IFNode(surrogate_function=surrogate.ATan()),
layer.SynapseFilter(tau=2., learnable=True),
layer.Linear(32, 10),
neuron.IFNode(surrogate_function=surrogate.ATan())
)
def __init__(self):
super().__init__()
self.fc = nn.Sequential(
layer.Linear(28, 32),
layer.LinearRecurrentContainer(
neuron.IFNode(surrogate_function=surrogate.ATan(), detach_reset=True),
in_features=32, out_features=32, bias=True
),
layer.Linear(32, 10),
neuron.IFNode(surrogate_function=surrogate.ATan())
)
def __init__(self, tau):
super().__init__()
self.layer = nn.Sequential(
layer.Flatten(),
layer.Linear(28 * 28, 10, bias=False),
neuron.LIFNode(tau=tau, surrogate_function=surrogate.ATan()),
)
def load_model(self, args, num_classes):
if args.model in spiking_resnet.__all__:
model = spiking_resnet.__dict__[args.model](
pretrained=args.pretrained,
spiking_neuron=neuron.IFNode,
surrogate_function=surrogate.ATan(),
detach_reset=True)
functional.set_step_mode(model, step_mode='m')
if args.cupy:
functional.set_backend(model, 'cupy', neuron.IFNode)
return model
else:
raise ValueError(
f"args.model should be one of {spiking_resnet.__all__}")
def main():
# python -m spikingjelly.activation_based.examples.classify_dvsg -T 16 -device cuda:0 -b 16 -epochs 64 -data-dir /datasets/DVSGesture/ -amp -cupy -opt adam -lr 0.001 -j 8
parser = argparse.ArgumentParser(description='Classify DVS Gesture')
parser.add_argument('-T',
default=16,
type=int,
help='simulating time-steps')
parser.add_argument('-device', default='cuda:0', help='device')
parser.add_argument('-b', default=16, type=int, help='batch size')
parser.add_argument('-epochs',
default=64,
type=int,
metavar='N',
help='number of total epochs to run')
parser.add_argument('-j',
default=4,
type=int,
metavar='N',
help='number of data loading workers (default: 4)')
parser.add_argument('-data-dir',
type=str,
help='root dir of DVS Gesture dataset')
parser.add_argument('-out-dir',
type=str,
default='./logs',
help='root dir for saving logs and checkpoint')
parser.add_argument('-resume',
type=str,
help='resume from the checkpoint path')
parser.add_argument('-amp',
action='store_true',
help='automatic mixed precision training')
parser.add_argument('-cupy', action='store_true', help='use cupy backend')
parser.add_argument('-opt',
type=str,
help='use which optimizer. SDG or Adam')
parser.add_argument('-momentum',
default=0.9,
type=float,
help='momentum for SGD')
parser.add_argument('-lr', default=0.1, type=float, help='learning rate')
parser.add_argument('-channels',
default=128,
type=int,
help='channels of CSNN')
args = parser.parse_args()
print(args)
net = parametric_lif_net.DVSGestureNet(channels=args.channels,
spiking_neuron=neuron.LIFNode,
surrogate_function=surrogate.ATan(),
detach_reset=True)
functional.set_step_mode(net, 'm')
if args.cupy:
functional.set_backend(net, 'cupy', instance=neuron.LIFNode)
print(net)
net.to(args.device)
train_set = DVS128Gesture(root=args.data_dir,
train=True,
data_type='frame',
frames_number=args.T,
split_by='number')
test_set = DVS128Gesture(root=args.data_dir,
train=False,
data_type='frame',
frames_number=args.T,
split_by='number')
train_data_loader = torch.utils.data.DataLoader(dataset=train_set,
batch_size=args.b,
shuffle=True,
drop_last=True,
num_workers=args.j,
pin_memory=True)
test_data_loader = torch.utils.data.DataLoader(dataset=test_set,
batch_size=args.b,
shuffle=True,
drop_last=False,
num_workers=args.j,
pin_memory=True)
scaler = None
if args.amp:
scaler = amp.GradScaler()
start_epoch = 0
max_test_acc = -1
optimizer = None
if args.opt == 'sgd':
optimizer = torch.optim.SGD(net.parameters(),
lr=args.lr,
momentum=args.momentum)
elif args.opt == 'adam':
optimizer = torch.optim.Adam(net.parameters(), lr=args.lr)
else:
raise NotImplementedError(args.opt)
lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
optimizer, args.epochs)
if args.resume:
checkpoint = torch.load(args.resume, map_location='cpu')
net.load_state_dict(checkpoint['net'])
optimizer.load_state_dict(checkpoint['optimizer'])
lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
start_epoch = checkpoint['epoch'] + 1
max_test_acc = checkpoint['max_test_acc']
out_dir = os.path.join(
args.out_dir,
f'T{args.T}_b{args.b}_{args.opt}_lr{args.lr}_c{args.channels}')
if args.amp:
out_dir += '_amp'
if args.cupy:
out_dir += '_cupy'
if not os.path.exists(out_dir):
os.makedirs(out_dir)
print(f'Mkdir {out_dir}.')
writer = SummaryWriter(out_dir, purge_step=start_epoch)
with open(os.path.join(out_dir, 'args.txt'), 'w',
encoding='utf-8') as args_txt:
args_txt.write(str(args))
args_txt.write('\n')
args_txt.write(' '.join(sys.argv))
for epoch in range(start_epoch, args.epochs):
start_time = time.time()
net.train()
train_loss = 0
train_acc = 0
train_samples = 0
for frame, label in train_data_loader:
optimizer.zero_grad()
frame = frame.to(args.device)
frame = frame.transpose(0, 1) # [N, T, C, H, W] -> [T, N, C, H, W]
label = label.to(args.device)
label_onehot = F.one_hot(label, 11).float()
if scaler is not None:
with amp.autocast():
out_fr = net(frame).mean(0)
loss = F.mse_loss(out_fr, label_onehot)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
else:
out_fr = net(frame).mean(0)
loss = F.mse_loss(out_fr, label_onehot)
loss.backward()
optimizer.step()
train_samples += label.numel()
train_loss += loss.item() * label.numel()
train_acc += (out_fr.argmax(1) == label).float().sum().item()
functional.reset_net(net)
train_time = time.time()
train_speed = train_samples / (train_time - start_time)
train_loss /= train_samples
train_acc /= train_samples
writer.add_scalar('train_loss', train_loss, epoch)
writer.add_scalar('train_acc', train_acc, epoch)
lr_scheduler.step()
net.eval()
test_loss = 0
test_acc = 0
test_samples = 0
with torch.no_grad():
for frame, label in test_data_loader:
frame = frame.to(args.device)
frame = frame.transpose(
0, 1) # [N, T, C, H, W] -> [T, N, C, H, W]
label = label.to(args.device)
label_onehot = F.one_hot(label, 11).float()
out_fr = net(frame).mean(0)
loss = F.mse_loss(out_fr, label_onehot)
test_samples += label.numel()
test_loss += loss.item() * label.numel()
test_acc += (out_fr.argmax(1) == label).float().sum().item()
functional.reset_net(net)
test_time = time.time()
test_speed = test_samples / (test_time - train_time)
test_loss /= test_samples
test_acc /= test_samples
writer.add_scalar('test_loss', test_loss, epoch)
writer.add_scalar('test_acc', test_acc, epoch)
save_max = False
if test_acc > max_test_acc:
max_test_acc = test_acc
save_max = True
checkpoint = {
'net': net.state_dict(),
'optimizer': optimizer.state_dict(),
'lr_scheduler': lr_scheduler.state_dict(),
'epoch': epoch,
'max_test_acc': max_test_acc
}
if save_max:
torch.save(checkpoint, os.path.join(out_dir, 'checkpoint_max.pth'))
torch.save(checkpoint, os.path.join(out_dir, 'checkpoint_latest.pth'))
print(args)
print(out_dir)
print(
f'epoch = {epoch}, train_loss ={train_loss: .4f}, train_acc ={train_acc: .4f}, test_loss ={test_loss: .4f}, test_acc ={test_acc: .4f}, max_test_acc ={max_test_acc: .4f}'
)
print(
f'train speed ={train_speed: .4f} images/s, test speed ={test_speed: .4f} images/s'
)
print(
f'escape time = {(datetime.datetime.now() + datetime.timedelta(seconds=(time.time() - start_time) * (args.epochs - epoch))).strftime("%Y-%m-%d %H:%M:%S")}\n'
)